instadeepai/Mava

[FEATURE] Create an `add_batch_dim` method

Open

#1,153 opened on 2025年1月9日

GitHub で見る
 (0 comments) (0 reactions) (0 assignees)Python (122 forks)auto 404
enhancementgood first issue

Repository metrics

Stars
 (922 stars)
PR merge metrics
 (PR metrics pending)

説明

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 を受け取り、`jax.tree.map(lambda x: x[jnp.newaxis], pytree)` を使用して新しいバッチ次元を追加した同じ pytree を返します。`jax.numpy as jnp` と `jax.tree util` をインポートしてください。
技術スタック
python
領域
backend
Issue 種別
機能
難度新しい貢献者にとっての実装難度です。1 は非常に小さな変更、5 は専門的な作業です。
1
推定時間経験ある貢献者が調査、実装、テスト、pull request 準備にかけるおおよその時間範囲です。
1時間未満
活動状況Issue が今どれだけ取り組みやすいかを示します。新しい、活発、古い、ブロック中、またはメンテナー入力待ちなどです。
新着
明確さ期待される変更、受け入れ条件、次の手順がどれだけ明確かを示します。
明確
前提条件
PythonJAX
初心者向け度初回貢献者にどれだけ取り組みやすいかを 1-100 で推定したスコアです。
90