CAX 项目使用教程

CAX 项目使用教程

1. 项目目录结构及介绍

CAX 项目是一个基于 JAX 的高性能、灵活的开源库，用于加速人工生命研究。项目目录结构如下：

• /.github/: 包含 GitHub 工作流文件，用于自动化任务，如代码测试和文档生成。
• /docs/: 存放项目文档的源文件。
• /examples/: 包含示例代码，用于展示如何使用 CAX 库。
• /src/: 项目的主要源代码目录，包含 CAX 库的实现。
  • /src/cax/: CAX 库的核心代码。
• /tests/: 包含用于测试 CAX 库的单元测试代码。
• /CONTRIBUTING.md: 提供贡献指南，描述如何为项目做贡献。
• /LICENSE: 项目使用的 MIT 许可证文件。
• /README.md: 项目的主页文档，概述项目的功能和用法。
• /pyproject.toml: Python 项目配置文件。
• /uv.lock: uv 工具的锁定文件，用于记录依赖项。

2. 项目的启动文件介绍

项目的启动通常通过 Python 脚本进行，但没有明确的单个启动文件。用户通常需要根据具体的使用场景编写入口脚本。以下是一个基本的 CAX 使用示例：

import jax
from flax import nnx
from cax.core.ca import CA
from cax.core.perceive import ConvPerceive
from cax.core.update import NCAUpdate

# 初始化参数
seed = 0
channel_size = 16
num_kernels = 3
hidden_layer_sizes = (128,)
cell_dropout_rate = 0.5
key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

# 创建感知和更新模块
perceive = ConvPerceive(
    channel_size=channel_size,
    perception_size=num_kernels * channel_size,
    rngs=rngs,
    feature_group_count=channel_size,
)
update = NCAUpdate(
    channel_size=channel_size,
    perception_size=num_kernels * channel_size,
    hidden_layer_sizes=hidden_layer_sizes,
    rngs=rngs,
    cell_dropout_rate=cell_dropout_rate,
    zeros_init=True,
)

# 创建 CA 实例
ca = CA(perceive, update)

# 初始化状态
state = jax.random.normal(key, (64, 64, channel_size))

# 运行 CA
state, metrics = ca(state, num_steps=128)

3. 项目的配置文件介绍

项目的配置主要通过 pyproject.toml 文件进行，该文件位于项目的根目录。这是一个示例配置：

[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"

[tool.setuptools]
packages = find:
python-requires = ">=3.10"

pyproject.toml 文件指定了构建系统所需的依赖，并告诉 setuptools 如何打包和分发项目。在这个文件中，你可以定义项目的依赖项和构建后端。

此外，CAX 项目可能使用环境变量和命令行参数来进行配置，具体取决于用户的使用方式和项目的具体需求。这些配置通常在项目的脚本或用户的环境中设置。