Index A | B | C | D | E | F | I | M | P | R | S | T | V | W A all_variants() (in module chex) assert_axis_dimension() (in module chex) assert_axis_dimension_comparator() (in module chex) assert_axis_dimension_gt() (in module chex) assert_axis_dimension_gteq() (in module chex) assert_axis_dimension_lt() (in module chex) assert_axis_dimension_lteq() (in module chex) assert_devices_available() (in module chex) assert_equal() (in module chex) assert_equal_rank() (in module chex) assert_equal_shape() (in module chex) assert_equal_shape_prefix() (in module chex) assert_equal_shape_suffix() (in module chex) assert_equal_size() (in module chex) assert_exactly_one_is_none() (in module chex) assert_gpu_available() (in module chex) assert_is_broadcastable() (in module chex) assert_is_divisible() (in module chex) assert_max_traces() (in module chex) assert_not_both_none() (in module chex) assert_numerical_grads() (in module chex) assert_rank() (in module chex) assert_scalar() (in module chex) assert_scalar_in() (in module chex) assert_scalar_negative() (in module chex) assert_scalar_non_negative() (in module chex) assert_scalar_positive() (in module chex) assert_shape() (in module chex) assert_size() (in module chex) assert_tpu_available() (in module chex) assert_tree_all_finite() (in module chex) assert_tree_has_only_ndarrays() (in module chex) assert_tree_is_on_device() (in module chex) assert_tree_is_on_host() (in module chex) assert_tree_is_sharded() (in module chex) assert_tree_no_nones() (in module chex) assert_tree_shape_prefix() (in module chex) assert_tree_shape_suffix() (in module chex) assert_trees_all_close() (in module chex) assert_trees_all_close_ulp() (in module chex) assert_trees_all_equal() (in module chex) assert_trees_all_equal_comparator() (in module chex) assert_trees_all_equal_dtypes() (in module chex) assert_trees_all_equal_shapes() (in module chex) assert_trees_all_equal_shapes_and_dtypes() (in module chex) assert_trees_all_equal_sizes() (in module chex) assert_trees_all_equal_structs() (in module chex) assert_type() (in module chex) B block_until_chexify_assertions_complete() (in module chex) C chexify() (in module chex) ChexVariantType (class in chex) clear_trace_counter() (in module chex) create_deprecated_function_alias() (in module chex) D dataclass() (in module chex) Dimensions (class in chex) disable_asserts() (in module chex) E enable_asserts() (in module chex) F fake_jit() (in module chex) fake_pmap() (in module chex) fake_pmap_and_jit() (in module chex) I if_args_not_none() (in module chex) M mappable_dataclass() (in module chex) P params_product() (in module chex) R register_dataclass_type_with_jax_tree_util() (in module chex) restrict_backends() (in module chex) S set_n_cpu_devices() (in module chex) T TestCase (class in chex) V variants() (in module chex) W warn_deprecated_function() (in module chex) warn_keyword_args_only_in_future() (in module chex) warn_only_n_pos_args_in_future() (in module chex) with_jittable_assertions() (in module chex)