Skip to content

Use shard in named_jit#178

Open
dlwh wants to merge 1 commit intomainfrom
codex/use-shard-with-axis-mapping-in-named_jit
Open

Use shard in named_jit#178
dlwh wants to merge 1 commit intomainfrom
codex/use-shard-with-axis-mapping-in-named_jit

Conversation

@dlwh
Copy link
Copy Markdown
Member

@dlwh dlwh commented Jul 21, 2025

Summary

  • remove use of out_shardings in named_jit
  • instead apply shard inside the jitted function

Testing

  • uv run pre-commit run --all-files
  • XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. uv run pytest tests

https://chatgpt.com/codex/tasks/task_e_687bef3b38f88331a41fd1d4d23950a6

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant