Skip to content

Add CLAUDE.md documenting decorator system and JAX jit boundary#223

Merged
Jammy2211 merged 1 commit intomainfrom
feature/jax_decorator_bypass
Mar 6, 2026
Merged

Add CLAUDE.md documenting decorator system and JAX jit boundary#223
Jammy2211 merged 1 commit intomainfrom
feature/jax_decorator_bypass

Conversation

@Jammy2211
Copy link
Owner

Summary

  • Adds CLAUDE.md documenting the core AbstractNDArray data structures, the @to_array / @to_grid / @to_vector_yx decorator system, the .array property pattern for accessing raw grid data inside function bodies, and why autoarray types cannot be returned from jax.jit.

Key content

  • Decorator system: how AbstractMaker selects the output wrapper type based on grid type (Grid2DArray2D, Grid2DIrregularArrayIrregular, etc.)
  • .array access pattern: inside decorated function bodies, use grid.array[:, 0] to get the raw numpy/jax array safely for both backends
  • JAX pytree constraint: autoarray types are not registered as JAX pytrees — they can be constructed inside a JIT trace but cannot be returned as JIT outputs
  • if xp is np: guard: the correct pattern for functions called directly inside jax.jit

🤖 Generated with Claude Code

Documents AbstractMaker, the to_array/to_grid/to_vector_yx decorators,
the .array property pattern for accessing raw grid data inside function
bodies, and why autoarray types cannot be returned from jax.jit.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Jammy2211 Jammy2211 merged commit 4ea1320 into main Mar 6, 2026
8 checks passed
@Jammy2211 Jammy2211 deleted the feature/jax_decorator_bypass branch March 6, 2026 16:15
Copy link
Collaborator

@rhayes777 rhayes777 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants