All pull requests

Pull requests

• Add test for interleaved constant and partial-input custom gradient captures.

  #288·copybara-service[bot] opened Sep 22, 2026·

  ·

• Support external custom gradient captures unused in forward outputs.

  #287·copybara-service[bot] opened Sep 22, 2026·

  ·

• Exclude internal captures from tf2jax gradient function arguments so they are not counted twice.

  #285·copybara-service[bot] opened Sep 21, 2026·

  ·

• 1. Pass recurse=True to _convert_all_gradient_functions inside builder() so nested custom gradients inside gradient functions are registered lazily on demand.

  #283·copybara-service[bot] opened Sep 18, 2026·

  ·

• Fix Copybara pattern for __init__.py pyrefly comment.

  #276·copybara-service[bot] opened Aug 21, 2026·

  ·

• Fix "AttributeError" thrown when processing multiple subgraphs induced by custom gradient definitions

  #268·nothingstopsme opened Apr 21, 2026·

  ·
  1

• Add codebase wiki

  #264·shaobohou opened Apr 12, 2026·

  ·
  1

• Upgrade GitHub Actions for Node 24 compatibility

  #262·salmanmkc opened Dec 16, 2025·

  ·

• Enhancements: Type annotations, CI/CD, error handling, docs, tests

  #257·genius-0963 opened Sep 7, 2025·

  ·

• Accelerate deprecation for attributes in jax.lib.xla_client and redirect dependencies to xla/python/xla_client.

  #247·copybara-service[bot] opened Jul 16, 2025·

  ·

• Improve Minor Errors for Clarity

  #242·yus100 opened May 9, 2025·

  ·

• [JAX] Deprecate jax.lib.xla_client.{heap_profile,mlir_api_version}.

  #234·copybara-service[bot] opened Apr 9, 2025·

  ·

• feat: add xlogy op

  #219·s22chan opened Oct 3, 2024·

  ·
  1

• fix(ops): sizes == -1 was interpreted incorrectly

  #218·s22chan opened Oct 3, 2024·

  ·
  2

• [JAX] add support for gather/scatter batching dims following the new attributes in stablehlo.

  #214·copybara-service[bot] opened Sep 21, 2024·

  ·

• OOB whatever; this CL description should be cleaned up before sent for review.

  #194·copybara-service[bot] opened May 28, 2024·

  ·
  1

• Implement "bad_indcies_policy" for ScatterNd.

  #193·copybara-service[bot] opened May 27, 2024·

  ·
  1

• ScatterNdFunctor to update operands of all valid indices and continue on bad indices.

  #192·copybara-service[bot] opened May 27, 2024·

  ·
  1

• Add "bad_indcies_policy" to ScatterNd

  #191·copybara-service[bot] opened May 27, 2024·

  ·
  1

• Bump tensorflow from 2.11.0 to 2.11.1

  dependencies
  #176·dependabot[bot] opened Dec 22, 2023·

  ·

• Use lax.cond for subsampled agent.

  #152·copybara-service[bot] opened Aug 9, 2023·

  ·
  1

• Temporary comment since I have tf2jax library changes on the diffbase.

  #151·copybara-service[bot] opened Aug 9, 2023·

  ·
  1