instadeepai/Mava

[FEATURE] Create an `add_batch_dim` method

Open

#1,153 创建于 2025年1月9日

在 GitHub 查看
 (0 评论) (0 反应) (0 负责人)Python (122 fork)auto 404
enhancementgood first issue

仓库指标

Star
 (922 star)
PR 合并指标
 (PR 指标待抓取)

描述

All over Mava we have statements like tree.map(lambda x: x[jnp.newaxis], pytree) I think it would make it a lot clearer if we had an add_batch_dim method in utils.jax_utils

贡献者指南

研究方向
在 `utils/jax utils.py` 中创建一个工具函数，该函数接受一个 pytree 并返回具有新批次维度的相同 pytree，通过 `jax.tree.map(lambda x: x[jnp.newaxis], pytree)` 实现。确保导入 `jax.numpy as jnp` 和 `jax.tree util`。
技术栈
python
领域
backend
议题类型
功能
难度面向新贡献者的预计实现难度，1 表示很小改动，5 表示专家级工作。
1
预计时间有经验贡献者完成调查、实现、测试并准备 pull request 的粗略时间范围。
1 小时以内
活动状态议题当前的可参与程度：新鲜、活跃、陈旧、阻塞或等待维护者输入。
新近可参与
清晰度议题是否清楚说明期望改动、验收标准和下一步。
清晰
前置要求
PythonJAX
新手友好度1-100 的估计分数，表示该议题对首次贡献者的友好程度。
90