Skip to content

Upgrading jax 0.11 to 0.12

Breaking0.11.2 → 0.12.0.1Breaking9Deprecated2New21Changed22Fixed31Other1

Full comparisonBack to jax
BreakingUpgrade verdict

9 breaking changes.

In plain terms
  • The old custom rematerialization function is gone; use the new custom_vjp-based approach instead.
  • Passing 2D vectors to cross now fails; the deprecation period has ended.
  • These functions now return tuples, matching NumPy 2.0 and later.
  • The shape you pass to multinomial no longer includes the outcome axis; drop its last entry.
Must act Top 5 of 11 that need code or config changes

What you must change

Breaking9
Removes jax.custom_remat in favor of jax.custom_vjp.defremat.
jax.custom_rematjax.custom_vjpjax.custom_vjp.defremat
0.12.0GitHub ReleaseSummarizedView source
jax.numpy.cross no longer accepts 2D vector inputs, completing a deprecation from JAX 0.11.0.
jax.numpy.cross
0.12.0GitHub ReleaseSummarizedView source
jax.numpy.atleast_1d, atleast_2d, and atleast_3d now return tuples of arrays instead of lists.
jax.numpy.atleast_1djax.numpy.atleast_2djax.numpy.atleast_3d
0.12.0GitHub ReleaseSummarizedView source
The shape argument of jax.random.multinomial is now the batch shape, excluding the outcome axis.
jax.random.multinomialshape
0.12.0GitHub ReleaseSummarizedView source
Under jax.vmap, custom_vjp functions with a primal output independent of batched inputs now raise an error instead of wrong gradients.
jax.vmapjax.custom_vjpgeneric_batching
0.12.0GitHub ReleaseSummarizedView source
The copy, order, and ndmin arguments to jax.numpy.array are now keyword-only.
jax.numpy.arraycopyorderndmin
0.12.0GitHub ReleaseSummarizedView source
jax.profiler.ProfileOptions now raises AttributeError when unknown attributes are assigned.
jax.profiler.ProfileOptionsAttributeError
0.12.0GitHub ReleaseSummarizedView source
jax.checkpoint, jax.custom_vjp, and jax.shard_map now error if the same Ref is passed more than once or both closed over and passed.
jax.checkpointjax.custom_vjpjax.shard_mapRef
0.12.0GitHub ReleaseSummarizedView source
Scalar constructors like jax.numpy.int32 and jax.numpy.float32 now have proper type annotations instead of Any, which may cause new typecheck errors.
jax.numpy.int32jax.numpy.float32jnp.dtype
0.12.0GitHub ReleaseSummarizedView source

Deprecations to plan for

Deprecated2
jax.lax.psend and jax.lax.precv are deprecated.
jax.lax.psendjax.lax.precv
0.12.0GitHub ReleaseSummarizedView source
Introduces jax.random.Key as an Array type alias; starting in JAX v0.13.0 it will be the return type for jax.random.key, split, and fold_in.
jax.random.Keyjax.random.keyjax.random.splitjax.random.fold_in
0.12.0GitHub ReleaseSummarizedView source

API changes

Generated by DepDelta from published Python source. Not mixed with prose notes.

  • removed: export custom_remat
  • removed: variable trace_ctx
  • removed: export TMABarrier
  • added: export CtaBarrier
  • added: export CtaBarrierRef
  • removed: class OptimizedTransferKind(enum.Enum)
  • removed: field OptimizedTransferKind.DOWNGRADABLE
  • removed: field OptimizedTransferKind.OPTIMIZED
  • removed: field OptimizedTransferKind.UNOPTIMIZED
  • changed: field IsTransferableSmemRegisters.optimized: OptimizedTransferKind → field IsTransferableSmemRegisters.optimized: bool
  • removed: class TMABarrier
  • removed: field TMABarrier.num_barriers: int
  • changed: function as_gpu_kernel(body, grid: tuple[int, int, int], block: tuple[int, int, int], in_shape, out_shape, smem_scratch_shape: ShapeTree | Union[ShapeTree], prof_spec: profiler.ProfilerSpec | None=…, cluster: tuple[int, int, int]=…, module_name: str=…, kernel_name: str | None=…, ir_version: int | None=…, thread_semantics: LoweringSemantics=…, inout_shape=…) → function as_gpu_kernel(body, grid: tuple[int, int, int], block: tuple[int, int, int], in_shape, out_shape, smem_scratch_shape: ShapeTree | Union[ShapeTree], prof_spec: profiler.ProfilerSpec | None=…, cluster: tuple[int, int, int]=…, module_name: str=…, kernel_name: str | None=…, ir_version: int | None=…, thread_semantics: LoweringSemantics=…, inout_shape=…, is_multi_process: bool=…)
  • added: class CtaBarrier
  • added: field CtaBarrier.arrival_count: int
  • added: field CtaBarrier.num_barriers: int
  • changed: field TiledLayout.lane_dims: tuple[int | Replicated, ...] → field TiledLayout.lane_dims: tuple[int | Replicated | Unreduced, ...]
  • changed: field TiledLayout.warp_dims: tuple[int | Replicated, ...] → field TiledLayout.warp_dims: tuple[int | Replicated | Unreduced, ...]
  • changed: method FragmentedArray.reduce(op: str | Callable[[ir.Value, ir.Value], ir.Value], axis: int | Sequence[int], scratch: ir.Value | None=…, *, acc_ilp: int | None=…) -> FragmentedArray → method FragmentedArray.reduce(op: str | Callable[[ir.Value, ir.Value], ir.Value], axis: int | Sequence[int], scratch: ir.Value | None=…, *, acc_ilp: int | None=…, target_layout: TiledLayout | None=…) -> FragmentedArray
  • changed: method TiledLayout.reduce(axes: Sequence[int]) -> TiledLayout → method TiledLayout.reduce(axes: Sequence[int], *, local_only: bool=…, op: str | None=…) -> TiledLayout
  • changed: method TiledLayout.remove_dimension(dim: int) -> TiledLayout → method TiledLayout.remove_dimension(dim: int, *, local_only: bool=…, op: str | None=…) -> TiledLayout
  • added: class Unreduced
  • added: field TiledLayout.unreduced_operation: str | None
  • added: field Unreduced.times: int
  • added: function plan_gmem_transfer(nested_ref_shape: Sequence[Sequence[int]], nested_ref_strides: Sequence[Sequence[int]], layout: TiledLayout, element_bits: int) -> TransferPlan
  • added: property FragmentedArray.is_unreduced -> bool
  • added: property TiledLayout.has_unreduced_dims -> bool
  • changed: field LaunchContext.tma_descriptors: dict[tuple[ir.Value, tuple[int, ...], int | None, tuple[MemRefTransform, ...], Any, int], ir.Value] → field LaunchContext.tma_descriptors: dict[Any, ir.Value]
  • changed: method LaunchContext.async_copy(*, src_ref: ir.Value, dst_ref: ir.Value, gmem_slice: Any=…, gmem_transform: MemRefTransform | tuple[MemRefTransform, ...]=…, gmem_peer_id: int | ir.Value | GlobalBroadcast | None=…, barrier: utils.BarrierRef | None=…, swizzle: int | None=…, arrive: bool | None=…, collective: Sequence[gpu.Dimension] | gpu.Dimension | None=…, leader_tracked: CopyPartition | None=…, p… → method LaunchContext.async_copy(*, src_ref: ir.Value, dst_ref: ir.Value, gmem_slice: Any=…, gmem_transform: MemRefTransform | tuple[MemRefTransform, ...]=…, gmem_peer_id: int | ir.Value | GlobalBroadcast | None=…, barrier: utils.BarrierRef | None=…, swizzle: int | None=…, arrive: bool | None=…, collective: tuple[gpu.Dimension, ...] | gpu.Dimension | None=…, leader_tracked: CopyPartition | None=…,…
  • changed: method LaunchContext.async_prefetch(*, gmem_ref: ir.Value, gmem_slice: Any=…, gmem_transform: MemRefTransform | tuple[MemRefTransform, ...]=…, gmem_peer_id: int | ir.Value | None=…, swizzle: int | None=…, collective: Sequence[gpu.Dimension] | gpu.Dimension | None=…, leader_tracked: CopyPartition | None=…, predicate: ir.Value | None | _DefaultPredicate=…) → method LaunchContext.async_prefetch(*, gmem_ref: ir.Value, gmem_slice: Any=…, gmem_transform: MemRefTransform | tuple[MemRefTransform, ...]=…, gmem_peer_id: int | ir.Value | None=…, swizzle: int | None=…, collective: tuple[gpu.Dimension, ...] | gpu.Dimension | None=…, leader_tracked: CopyPartition | None=…, predicate: ir.Value | None | _DefaultPredicate=…)
  • added: method LaunchContext.async_copy_im2col(*, src_ref: ir.Value, dst_ref: ir.Value, barrier: utils.BarrierRef, window_start_bounds: Sequence[tuple[int, int]], start_indices: Sequence[int | ir.Value], filter_offsets: Sequence[int | ir.Value], window_strides: Sequence[int] | None=…, swizzle: int | None=…, arrive: bool=…, predicate: ir.Value | None | _DefaultPredicate=…)
  • added: variable HBM_ALIGNMENT
  • added: field DerivationContext.op_for_constraint: dict[cs.Constraint, ir.OpView]
  • added: function diagnose_unsatisfiable_system(ctx: DerivationContext, initial_system: cs.ConstraintSystem)
  • added: function try_raise_constraint_specific_error(op: ir.OpView, constraint: cs.Constraint)
  • added: function pprint_layout(layout: fa.FragmentedLayout) -> str
  • changed: function commit_arrive(barrier: utils.BarrierRef | ir.Value, collective: bool=…, ctx: LaunchContext | None=…) -> None → function commit_arrive(barrier: utils.BarrierRef | ir.Value, collective: bool=…, ctx: LaunchContext | None=…, *, predicate: ir.Value | None=…) -> None
  • changed: function parse_indices(index, shape: Sequence[int], *, check_oob: bool=…) -> tuple[list[ir.Value | int], list[int], list[bool]] → function parse_indices(index, shape: Sequence[int], *, check_oob: bool=…, allow_negative: bool=…) -> tuple[list[ir.Value | int], list[int], list[bool]]
  • changed: method BarrierRef.arrive_expect_tx(tx_count: int | ir.Value, predicate: ir.Value | None=…) → method BarrierRef.arrive_expect_tx(tx_count: int | ir.Value, predicate: ir.Value | None=…, tensor_core_order_scope: ThreadSubset | None=…)
  • added: class CtaBarrierRef
  • added: field CtaBarrierRef.arrival_count: int
  • added: field CtaBarrierRef.base_id: int
  • added: field CtaBarrierRef.num_barriers: int
  • added: field CtaBarrierRef.offset: ir.Value
  • added: function before_thread_sync(*, sync_threads: bool, scope: ThreadSubset=…)
  • added: function memref_descriptor_type(memref_ty: ir.MemRefType) -> llvm.StructType
  • added: function to_memref_descriptor(memref_arg: ir.Value) -> ir.Value
  • added: method CtaBarrierRef.arrive(*, predicate: ir.Value | None=…) -> None
  • added: method CtaBarrierRef.arrive_and_wait() -> None
  • added: method CtaBarrierRef.get_barrier_id() -> ir.Value
  • added: static method CtaBarrierRef.initialize(base_id: int, arrival_count: int=…, num_barriers: int=…) -> CtaBarrierRef
  • removed: function control_dep(src, dst)
  • removed: function schedule(ops)
  • removed: export DeviceIdType
  • added: variable DeviceIdType
  • changed: function get_neighbor(idx: jax.Array, mesh: jax.sharding.Mesh, axis_name: str, *, direction: str) -> tuple[jax.Array, ...] → function get_neighbor(idx: jax.Array, mesh: jax.sharding.Mesh, axis_name: str, *, direction: str) -> dict[str, jax.Array]
  • changed: function flash_attention_kernel(data_next_ref, block_mask_ref, mask_next_ref, q_ref, k_ref, v_ref, q_segment_ids_ref, kv_segment_ids_ref, sinks_ref, mask_ref, q_sequence_ref, m_scratch_ref, l_scratch_ref, o_scratch_ref, o_ref, logsumexp_ref=…, *, mask_value: float, grid_width: int, bq: int, bkv: int, bkv_compute: int, head_dim_v: int, q_layout: QKVLayout, k_layout: QKVLayout, v_layout: QKVLayout,… → function flash_attention_kernel(data_next_ref, block_mask_ref, mask_next_ref, q_ref, k_ref, v_ref, q_segment_ids_ref, kv_segment_ids_ref, sinks_ref, mask_ref, q_sequence_ref, o_ref, logsumexp_ref, m_scratch_ref, l_scratch_ref, o_scratch_ref, *, mask_value: float, grid_width: int, bq: int, bkv: int, bkv_compute: int, head_dim_v: int, q_layout: QKVLayout, k_layout: QKVLayout, v_layout: QKVLayout, a…
  • added: variable BufferedRefBase
  • removed: export precv
  • removed: export psend
  • added: export polynomial
  • added: export polynomial_p
  • added: export psum_scatter_start
  • added: variable precv
  • added: variable psend
  • changed: function get_scaled_dot_general_config(mode: Literal['nvfp4', 'mxfp8'], global_scale: Array | None=…) -> BlockScaleConfig → function get_scaled_dot_general_config(mode: Literal['nvfp4', 'mxfp8', 'mxfp4'], global_scale: Array | None=…) -> BlockScaleConfig
  • changed: function atleast_1d() -> list[Array] | (x: ArrayLike, /) -> Array | (x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array] → function atleast_1d() -> tuple[()] | (x: ArrayLike, /) -> Array | (x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> tuple[Array, ...]
  • changed: function atleast_2d() -> list[Array] | (x: ArrayLike, /) -> Array | (x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array] → function atleast_2d() -> tuple[()] | (x: ArrayLike, /) -> Array | (x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> tuple[Array, ...]
  • changed: function atleast_3d() -> list[Array] | (x: ArrayLike, /) -> Array | (x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array] → function atleast_3d() -> tuple[()] | (x: ArrayLike, /) -> Array | (x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> tuple[Array, ...]
  • changed: variable bfloat16: Any → variable bfloat16: _ScalarType
  • changed: variable bool: Any → variable bool: _ScalarType
  • changed: variable bool_: Any → variable bool_: _ScalarType
  • changed: variable cdouble: Any → variable cdouble: _ScalarType
  • changed: variable complex128: Any → variable complex128: _ScalarType
  • changed: variable complex64: Any → variable complex64: _ScalarType
  • changed: variable complex_: Any → variable complex_: _ScalarType
  • changed: variable csingle: Any → variable csingle: _ScalarType
  • changed: variable double: Any → variable double: _ScalarType
  • changed: variable float16: Any → variable float16: _ScalarType
  • changed: variable float32: Any → variable float32: _ScalarType
  • changed: variable float4_e2m1fn: Any → variable float4_e2m1fn: _ScalarType
  • changed: variable float64: Any → variable float64: _ScalarType
  • changed: variable float8_e3m4: Any → variable float8_e3m4: _ScalarType
  • changed: variable float8_e4m3: Any → variable float8_e4m3: _ScalarType
  • changed: variable float8_e4m3b11fnuz: Any → variable float8_e4m3b11fnuz: _ScalarType
  • changed: variable float8_e4m3fn: Any → variable float8_e4m3fn: _ScalarType
  • changed: variable float8_e4m3fnuz: Any → variable float8_e4m3fnuz: _ScalarType
  • changed: variable float8_e5m2: Any → variable float8_e5m2: _ScalarType
  • changed: variable float8_e5m2fnuz: Any → variable float8_e5m2fnuz: _ScalarType
  • changed: variable float8_e8m0fnu: Any → variable float8_e8m0fnu: _ScalarType
  • changed: variable float_: Any → variable float_: _ScalarType
  • changed: variable int16: Any → variable int16: _ScalarType
  • changed: variable int2: Any → variable int2: _ScalarType
  • changed: variable int32: Any → variable int32: _ScalarType
  • changed: variable int4: Any → variable int4: _ScalarType
  • changed: variable int64: Any → variable int64: _ScalarType
  • changed: variable int8: Any → variable int8: _ScalarType
  • changed: variable int_: Any → variable int_: _ScalarType
  • changed: variable single: Any → variable single: _ScalarType
  • changed: variable uint16: Any → variable uint16: _ScalarType
  • changed: variable uint2: Any → variable uint2: _ScalarType
  • changed: variable uint32: Any → variable uint32: _ScalarType
  • changed: variable uint4: Any → variable uint4: _ScalarType
  • changed: variable uint64: Any → variable uint64: _ScalarType
  • changed: variable uint8: Any → variable uint8: _ScalarType
  • changed: variable uint: Any → variable uint: _ScalarType
  • added: variable float6_e2m3fn: _ScalarType
  • added: variable float6_e3m2fn: _ScalarType