Skip to content

Use jax.Array for JAX functions. - #2459

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_981040772
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_981040772

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Sep 14, 2026 •

Copy link
Copy Markdown

Use jax.Array for JAX functions.

@copybara-service copybara-service Bot changed the title Move stencil JIT out of _compute_inner_grad. Use jax.Array for JAX functions. Sep 14, 2026
PiperOrigin-RevId: 981040772
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.

1 participant