jax-ml/jax

Pattern Findings

11935 matches across 19 categories. Click a row to expand file-level details.

Hyper-Verbose Identifiers5566 hits · 4365 pts

Severity	File	Line	Snippet
LOW	ci/parse_wheel_metadata.py	34	def _extract_wheel_package_name(wheel_name: str) -> str \| None:
LOW	ci/parse_wheel_metadata.py	52	def parse_expected_wheel_versions(wheel_dir: Path) -> dict[str, str]:
LOW	tests/random_test.py	569	def test_key_construction_with_dtype(self, dtype_spec):
LOW	tests/random_test.py	573	def test_key_construction_with_both_impl_and_dtype(self):
LOW	tests/random_test.py	583	def test_wrap_key_data_with_dtype(self, dtype_spec):
LOW	tests/random_test.py	591	def test_wrap_key_data_with_both_impl_and_dtype(self):
LOW	tests/random_test.py	1215	def test_key_make_like_other_key_via_dtype(self, prng_name):
LOW	tests/random_test.py	1223	def test_key_wrap_like_other_key_via_dtype(self, prng_name):
LOW	tests/random_test.py	1231	def test_key_impl_from_string_error(self):
LOW	tests/random_test.py	1235	def test_key_impl_from_object_error(self):
LOW	tests/random_test.py	1242	def test_key_impl_builtin_is_string_name(self, name):
LOW	tests/random_test.py	199	def test_config_prngs_registered(self):
LOW	tests/random_test.py	412	def test_threefry_gpu_kernel_lowering(self):
LOW	tests/random_test.py	445	def test_threefry_split_fold_in_symmetry(self, make_key):
LOW	tests/random_test.py	458	def test_threefry_split_vmapped_fold_in_symmetry(self, make_key):
LOW	tests/random_test.py	472	def test_loggamma_nan_corner_case(self):
LOW	tests/random_test.py	545	def test_default_prng_selection(self, make_key, name, impl):
LOW	tests/random_test.py	557	def test_key_construction_with_explicit_impl_name(self, make_key, name, impl):
LOW	tests/random_test.py	609	def test_legacy_prng_key_flag(self):
LOW	tests/random_test.py	637	def test_seed_no_implicit_transfers(self, make_key):
LOW	tests/random_test.py	713	def test_construction_upgrade_flag(self):
LOW	tests/random_test.py	759	def test_key_dtype_attributes(self):
LOW	tests/random_test.py	814	def test_cpp_dispatch_aot_normal(self):
LOW	tests/random_test.py	827	def test_cpp_dispatch_aot_split(self):
LOW	tests/random_test.py	905	def test_eval_shape_keys_in_out(self):
LOW	tests/random_test.py	963	def test_dynamic_update_slice(self):
LOW	tests/random_test.py	1053	def test_device_put_replicated(self):
LOW	tests/random_test.py	1059	def test_make_array_from_callback(self):
LOW	tests/random_test.py	1071	def test_make_array_from_single_device_arrays(self):
LOW	tests/random_test.py	1081	def test_key_array_custom_jvp(self):
LOW	tests/random_test.py	1111	def test_key_array_indexing_nd(self):
LOW	tests/random_test.py	1136	def test_array_impl_attributes(self):
LOW	tests/random_test.py	1194	def test_key_make_like_other_key(self, prng_name):
LOW	tests/random_test.py	1204	def test_key_wrap_like_other_key(self, prng_name):
LOW	tests/random_test.py	1262	def test_keyarray_custom_vjp_symbolic_zeros(self):
LOW	tests/random_test.py	1277	def test_keyarray_array_conversion_fails(self):
LOW	tests/random_test.py	1301	def _double_threefry_random_bits(key, bit_width, shape):
LOW	tests/random_test.py	1565	def test_full_like_with_key_fillvalue(self):
LOW	tests/random_test.py	1586	def test_full_with_key_fillvalue(self):
LOW	tests/extend_test.py	84	def test_key_make_with_custom_impl(self):
LOW	tests/extend_test.py	92	def test_key_wrap_with_custom_impl(self):
LOW	tests/extend_test.py	102	def test_key_make_with_custom_impl_via_dtype(self):
LOW	tests/extend_test.py	109	def test_key_wrap_with_custom_impl_via_dtype(self):
LOW	tests/extend_test.py	118	def test_key_dtype_and_spec_with_custom_impl(self):
LOW	tests/extend_test.py	130	def test_unknown_platform_error(self):
LOW	tests/extend_test.py	139	def test_hlo_sharding_roundtrip(self):
LOW	tests/scheduling_groups_test.py	67	def test_xla_metadata_call_inlineable(self):
LOW	tests/scheduling_groups_test.py	88	def test_xla_metadata_call_inlineable_remat_in_scan(self):
LOW	tests/scheduling_groups_test.py	105	def test_xla_metadata_call_deduplication(self):
LOW	tests/scheduling_groups_test.py	133	def test_xla_metadata_call_deduplication_remat(self):
LOW	tests/scheduling_groups_test.py	162	def test_xla_metadata_call_deduplication_kwargs(self):
LOW	tests/jax_to_ir_test.py	65	def test_jax_to_hlo_with_constants(self):
LOW	tests/jax_to_ir_test.py	85	def test_parse_shape_str_invalid(self):
LOW	tests/ragged_collective_test.py	133	def test_ragged_all_to_all_grad(self, axis_name, mesh_axes):
LOW	tests/ragged_collective_test.py	207	def test_ragged_all_to_all_axis_index_groups(self, axis_name, mesh_axes):
LOW	tests/ragged_collective_test.py	301	def test_ragged_all_to_all_degenerate_groups(self, axis_name, mesh_axes):
LOW	tests/ragged_collective_test.py	384	def test_ragged_all_to_all_vmap_multi_dim_operand(self):
LOW	tests/ragged_collective_test.py	480	def test_ragged_all_to_all_vmap(
LOW	tests/ragged_collective_test.py	581	def test_ragged_all_to_all_vmap_unsupported_axis_index_groups(self):
LOW	tests/ragged_collective_test.py	639	def test_ragged_all_to_all_errors(self):
5506 more matches not shown…

Unused Imports3476 hits · 2452 pts

Severity	File	Line
LOW	tests/random_test.py	15
LOW	tests/batching_test.py	15
LOW	tests/debug_info_test.py	15
LOW	tests/mutable_array_test.py	15
LOW	tests/lax_numpy_test.py	15
LOW	tests/shard_map_test.py	15
LOW	tests/export_harnesses_multi_platform_test.py	22
LOW	tests/shape_poly_test.py	15
LOW	tests/lax_vmap_test.py	15
LOW	tests/lax_numpy_indexing_test.py	15
LOW	tests/typing_test.py	21
LOW	tests/array_api_test.py	20
LOW	tests/state_test.py	15
LOW	tests/export_test.py	14
LOW	tests/export_test.py	46
LOW	tests/export_test.py	53
LOW	tests/pmap_test.py	15
LOW	tests/colocated_python_test.py	36
LOW	tests/lru_cache_test.py	15
LOW	tests/compilation_cache_test.py	15
LOW	tests/hijax_test.py	15
LOW	tests/export_serialization_back_compat_test.py	53
LOW	tests/api_test.py	15
LOW	tests/roofline_test.py	14
LOW	tests/lax_test.py	14
LOW	tests/mosaic/gpu_torch_test.py	29
LOW	tests/mosaic/flash_attention_test.py	24
LOW	tests/mosaic/gpu_test.py	55
LOW	tests/multiprocess/colocated_python_test.py	24
LOW	tests/pallas/tpu_all_gather_test.py	16
LOW	tests/pallas/tpu_splash_attention_mask_test.py	15
LOW	tests/pallas/tpu_pallas_pipeline_test.py	16
LOW	tests/pallas/pallas_shape_poly_test.py	31
LOW	tests/pallas/tpu_pallas_random_test.py	26
LOW	tests/pallas/indexing_test.py	15
LOW	tests/pallas/tpu_splash_attention_kernel_test.py	15
LOW	tests/pallas/pallas_test.py	14
LOW	docs/parallel.py	61
LOW	docs/autodidax.py	2147
LOW	jaxlib/xla_client.py	17
LOW	jaxlib/init.py	15
LOW	jaxlib/mosaic/gpu/wheel/__init__.py	17
LOW	jaxlib/mosaic/python/tpu.py	20
LOW	jaxlib/mosaic/python/tpu.py	22
LOW	jaxlib/mosaic/python/tpu.py	24
LOW	jaxlib/mosaic/python/mosaic_gpu.py	23
LOW	jaxlib/mosaic/python/mosaic_gpu.py	25
LOW	jaxlib/mosaic/python/mosaic_gpu.py	26
LOW	jaxlib/tools/build_utils.py	17
LOW	jaxlib/triton/dialect.py	19
LOW	jaxlib/triton/dialect.py	23
LOW	jaxlib/triton/dialect.py	23
LOW	jaxlib/triton/dialect.py	30
LOW	jaxlib/triton/dialect.py	31
LOW	jax/sharding.py	18
LOW	jax/sharding.py	19
LOW	jax/sharding.py	19
LOW	jax/sharding.py	19
LOW	jax/sharding.py	19
LOW	jax/sharding.py	19
3416 more matches not shown…

Docstring Block Structure566 hits · 2302 pts

Severity	File	Line	Snippet
HIGH	jax/experimental/multihost_utils.py	549	A context manager for atomically running code on the set of live devices. THIS API IS UNDER ACTIVE DEVELOPMENT AND IS
HIGH	jax/experimental/mosaic/gpu/profiler.py	142	Measures the GPU runtime of a function using CUPTI. ``measure`` is a higher-order function that wraps a function ``f`
HIGH	…ops/tpu/splash_attention/splash_attention_mask_info.py	84	Downcast numpy array. If possible, downcast the data-type of the input array to the smallest numpy type (among np.i
HIGH	…ops/tpu/splash_attention/splash_attention_mask_info.py	326	Similar to `_process_mask` but the mask must be a dynamic array. Since the mask is dynamic, we can't know the exact n
HIGH	…ops/tpu/splash_attention/splash_attention_mask_info.py	528	Transform a dense mask into a sparse representation. The number of head and Q sequence shards are needed to create a
HIGH	…llas/ops/tpu/splash_attention/splash_attention_mask.py	97	Makes a chunked causal attention mask. Args: shape: The desired shape of the mask (q_seq_len, kv_seq_len). ch
HIGH	jax/experimental/sparse/linalg.py	43	Compute the top-k standard eigenvalues using the LOBPCG routine. LOBPCG [1] stands for Locally Optimal Block Precondi
HIGH	jax/experimental/sparse/bcoo.py	2437	Experimental batched COO matrix implemented in JAX Args: (data, indices) : data and indices in batched COO format
HIGH	jax/_src/callback.py	269	Calls a pure Python callback. Works under :func:`jit`/:func:`~vmap`/etc. For more explanation, see `External Callback
HIGH	jax/_src/tree.py	26	Call all() over the leaves of a tree. Args: tree: the pytree to evaluate is_leaf : an optionally specified fu
HIGH	jax/_src/tree.py	55	Flattens a pytree. The flattening order (i.e. the order of elements in the output list) is deterministic, correspon
HIGH	jax/_src/tree.py	91	Gets the leaves of a pytree. Args: tree: the pytree for which to get the leaves is_leaf : an optionally speci
HIGH	jax/_src/tree.py	120	Maps a multi-input function over pytree args to produce a new pytree. Args: f: function that takes ``1 + len(rest
HIGH	jax/_src/tree.py	163	Call reduce() over the leaves of a tree. Args: function: the reduction function tree: the pytree to reduce ov
HIGH	jax/_src/tree.py	203	Perform a reduction over a pytree with an associative binary operation. This function exploits the fact that the oper
HIGH	jax/_src/tree.py	245	Gets the treedef for a pytree. Args: tree: the pytree for which to get the leaves is_leaf : an optionally spe
HIGH	jax/_src/tree.py	273	Transform a tree having tree structure (outer, inner) into one having structure (inner, outer). Args: outer_treed
HIGH	jax/_src/tree.py	303	Reconstructs a pytree from the treedef and the leaves. The inverse of :func:`tree_flatten`. Args: treedef: the
HIGH	jax/_src/tree.py	335	Flattens a pytree like ``tree_flatten``, but also returns each leaf's key path. Args: tree: a pytree to flatten.
HIGH	jax/_src/tree.py	366	Gets the leaves of a pytree like ``tree_leaves`` and returns each leaf's key path. Args: tree: a pytree. If it co
HIGH	jax/_src/tree.py	395	Maps a multi-input function over pytree key path and args to produce a new pytree. This is a more powerful alternativ
HIGH	jax/_src/tree.py	433	Broadcasts a tree prefix into the full structure of a given tree. Args: prefix_tree: a pytree that is a tree
HIGH	jax/_src/xla_bridge.py	317	Creates a CPU client with the requested collectives implementation. The implementation of CPU collectives used by the
HIGH	jax/_src/config.py	368	Set up thread-local state and return a contextmanager for managing it. This function is a convenience wrapper. It def
HIGH	jax/_src/core.py	2891	Invalidate a given reference and return its final value. For more information about mutable array references, refer t
HIGH	jax/_src/dtypes.py	887	Returns the type to which a binary operation should cast its arguments. JAX implementation of :func:`numpy.promote_ty
HIGH	jax/_src/dtypes.py	1135	Check if a dtype/value is safe to cast to another dtype/value Args: input_dtype_or_value: a dtype or value (to be
HIGH	jax/_src/api.py	208	Sets up ``fun`` for just-in-time compilation with XLA. Args: fun: Function to be jitted. ``fun`` should be a pure
HIGH	jax/_src/api.py	2203	Transfer array shards to specified devices and form Array(s). Args: shards: A sequence of arrays, scalars, or (ne
HIGH	jax/_src/api.py	2283	Transfer array(s) to each specified device and form Array(s). Args: x: an array, scalar, or (nested) standard Pyt
HIGH	jax/_src/api.py	2352	Transfer ``x`` to host. If ``x`` is a pytree, then the individual buffers are copied in parallel. Args: x: An
HIGH	jax/_src/api.py	2502	A context manager that adds a user specified name to the JAX name stack. When staging out computations for just-in-ti
HIGH	jax/_src/distributed.py	233	Initializes the JAX distributed system. Calling :func:`~jax.distributed.initialize` prepares JAX for execution on m
HIGH	jax/_src/sharding_impls.py	650	Computes the global shape given the per process if possible. The returned shape will have the size of the global tens
HIGH	jax/_src/tree_util.py	125	Makes a tuple treedef from an iterable of child treedefs. Args: treedefs: iterable of PyTree structures Return
HIGH	jax/_src/tree_util.py	151	Return a list of treedefs for immediate children Args: treedef: a single PyTreeDef Returns: a list of PyTr
HIGH	jax/_src/tree_util.py	176	Return True if the treedef represents a leaf. Args: treedef: tree to check Returns: True if treedef is a l
HIGH	jax/_src/tree_util.py	204	Tests whether all elements in the given iterable are all leaves. This function is useful in advanced cases, for examp
HIGH	jax/_src/tree_util.py	336	Extends the set of types that are considered internal nodes in pytrees. This function is a thin wrapper around ``regi
HIGH	jax/_src/tree_util.py	688	Flatten the given pytree node by one level. Args: tree: A valid pytree node, either built-in or registered via
HIGH	jax/_src/tree_util.py	834	Helper to pretty-print a tuple of keys. Args: keys: A tuple of ``KeyEntry`` or any class that can be converted to
HIGH	jax/_src/tree_util.py	951	Extends the set of types that are considered internal nodes in pytrees. This function is similar to ``register_pytree
HIGH	jax/_src/tree_util.py	1005	Extends the set of types that are considered internal nodes in pytrees. This differs from ``register_pytree_with_keys
HIGH	jax/_src/tree_util.py	1191	Registers `cls` as a pytree with no leaves. Instances are treated as static by :func:`jax.jit`, :func:`jax.pmap`, etc
HIGH	jax/_src/custom_partitioning_sharding_rule.py	215	Parses the LHS or RHS of an Einsum notation like string. Converts each operand or result in the Einsum notation like
HIGH	jax/_src/array.py	691	Returns a ``jax.Array`` via data fetched from ``data_callback``. ``data_callback`` is used to fetch the data for each
HIGH	jax/_src/array.py	826	Creates distributed tensor using the data available in process. This function is a common special case of `make_array
HIGH	jax/_src/array.py	1026	Returns a ``jax.Array`` from a sequence of ``jax.Array``\s each on a single device. Every device in input ``shardi
HIGH	jax/_src/profiler.py	94	Registers a subprocess's profiler server to be profiled alongside the current process. When the current process colle
HIGH	jax/_src/mesh_utils.py	789	Creates a performant device mesh for jax.sharding.Mesh. Args: mesh_shape: shape of logical mesh, ordered by incre
HIGH	jax/_src/mesh_utils.py	870	Creates a device mesh for hybrid (e.g., ICI and DCN) parallelism. Args: mesh_shape: shape of the logical mesh for
HIGH	jax/_src/pmap.py	477	Infer axis size from the first mapped argument. shard_map already does a check on all arguments, so just look at firs
HIGH	jax/_src/pmap.py	528	Extract dynamic args and argnums after handling static args. Args: wrapped_f: The wrapped function. static_br
HIGH	jax/_src/pmap.py	566	Compute flat in_axes tuple from in_axes prefix and args structure. Args: in_axes: The original in_axes specificat
HIGH	jax/_src/pmap.py	673	Compute effective mesh devices based on context. Args: devices: The mesh devices tuple. backend: The backend
HIGH	jax/_src/custom_derivatives.py	190	Define a custom JVP rule for the function represented by this instance. Args: jvp: a Python callable represen
HIGH	jax/_src/custom_derivatives.py	239	Convenience wrapper for defining JVPs for each argument separately. This convenience wrapper cannot be used togethe
HIGH	jax/_src/custom_derivatives.py	624	Define a custom VJP rule for the function represented by this instance. Args: fwd: a Python callable represen
HIGH	jax/_src/interpreters/mlir.py	259	Translate a Python ``val`` to an IR constant. See https://docs.jax.dev/en/latest/internals/constants.html. Args:
HIGH	jax/_src/interpreters/mlir.py	291	Translate a Python ``val`` to a sequence of IR constants. See https://docs.jax.dev/en/latest/internals/constants.html
506 more matches not shown…

Over-Commented Block1399 hits · 1334 pts

Severity	File	Line	Snippet
LOW	conftest.py	1	# Copyright 2021 The JAX Authors.
LOW	conftest.py	41	# The pytest_collection hook can be used to overwrite the collection logic, but
LOW	setup.py	1	# Copyright 2018 The JAX Authors.
LOW	build_wheel.py	1	# Copyright 2025 The JAX Authors.
LOW	ci/run_pytest_tpu.sh	1	#!/bin/bash
LOW	ci/run_pytest_rocm.sh	1	#!/bin/bash
LOW	ci/run_bazel_test_cuda_non_rbe.sh	1	#!/bin/bash
LOW	ci/run_bazel_test_cuda_non_rbe.sh	21	# -e: abort script if one command fails
LOW	ci/run_bazel_cuda_targeted_tests.sh	1	#!/bin/bash
LOW	ci/run_bazel_cuda_targeted_tests.sh	21	# JAXCI_BAZEL_TARGETS: newline-separated Bazel targets.
LOW	ci/run_bazel_test_rocm_rbe.sh	1	#!/bin/bash
LOW	ci/build_artifacts.sh	1	#!/bin/bash
LOW	ci/build_rocm_artifacts.sh	1	#!/bin/bash
LOW	ci/run_bazel_test_cpu_rbe.sh	1	#!/bin/bash
LOW	ci/upload_rocm_logs.sh	1	#!/bin/bash
LOW	ci/run_bazel_test_cuda_rbe.sh	1	#!/bin/bash
LOW	ci/run_pytest_cuda.sh	1	#!/bin/bash
LOW	ci/parse_wheel_metadata.py	1	# Copyright 2026 The JAX Authors.
LOW	ci/run_bazel_test_tpu.sh	1	#!/bin/bash
LOW	ci/run_bazel_test_tpu.sh	181	--test_env=JAX_PLATFORMS=tpu,cpu \
LOW	ci/run_pytest_cpu.sh	1	#!/bin/bash
LOW	ci/utilities/install_wheels_locally.sh	1	#!/bin/bash
LOW	ci/utilities/run_docker_container.sh	1	#!/bin/bash
LOW	ci/utilities/run_docker_container.sh	21	# Note: While GitHub action workflows use the same Docker images, they do not
LOW	ci/utilities/convert_msys_paths_to_win_paths.py	1	# Copyright 2024 The JAX Authors.
LOW	ci/utilities/collect_bazel_test_xmls.sh	1	#!/bin/bash
LOW	ci/utilities/collect_bazel_test_xmls.sh	21	# bazel-testlogs/tests/cpu_tests/test.xml). Pytest, by contrast, writes a
LOW	ci/utilities/setup_build_environment.sh	1	#!/bin/bash
LOW	ci/utilities/run_auditwheel.sh	1	#!/bin/bash
LOW	ci/postprocess/process_test_results.sh	1	#!/bin/bash
LOW	ci/postprocess/xml2json.py	1	# Copyright 2026 The JAX Authors.
LOW	tests/random_test.py	1	# Copyright 2018 The JAX Authors.
LOW	tests/random_test.py	641
LOW	tests/extend_test.py	1	# Copyright 2023 The JAX Authors.
LOW	tests/svd_test.py	1	# Copyright 2022 The JAX Authors.
LOW	tests/scipy_spatial_test.py	1	# Copyright 2023 The JAX Authors.
LOW	tests/scheduling_groups_test.py	1	# Copyright 2025 The JAX Authors.
LOW	tests/jax_to_ir_test.py	1	# Copyright 2019 The JAX Authors.
LOW	tests/ragged_collective_test.py	1	# Copyright 2025 The JAX Authors.
LOW	tests/pytorch_interoperability_test.py	1	# Copyright 2020 The JAX Authors.
LOW	tests/fft_test.py	1	# Copyright 2019 The JAX Authors.
LOW	tests/export_back_compat_test.py	1	# Copyright 2023 The JAX Authors.
LOW	tests/export_back_compat_test.py	1041	# stablehlo.dynamic_rbg_bit_generator is used temporarily for a
LOW	tests/mesh_utils_test.py	1	# Copyright 2021 The JAX Authors. All Rights Reserved.
LOW	tests/lax_vmap_op_test.py	1	# Copyright 2020 The JAX Authors.
LOW	tests/debug_nans_test.py	1	# Copyright 2019 The JAX Authors.
LOW	tests/lobpcg_test.py	1	# Copyright 2022 The JAX Authors.
LOW	tests/heap_profiler_test.py	1	# Copyright 2021 The JAX Authors.
LOW	tests/custom_root_test.py	1	# Copyright 2018 The JAX Authors.
LOW	tests/mosaic_test.py	1	# Copyright 2023 The JAX Authors.
LOW	tests/lax_control_flow_test.py	1	# Copyright 2018 The JAX Authors.
LOW	tests/lax_control_flow_test.py	2221
LOW	tests/array_interoperability_test.py	1	# Copyright 2020 The JAX Authors.
LOW	tests/lax_scipy_test.py	1	# Copyright 2018 The JAX Authors.
LOW	tests/errors_test.py	1	# Copyright 2020 The JAX Authors.
LOW	tests/distributed_initialize_test.py	1	# Copyright 2025 The JAX Authors.
LOW	tests/debugger_test.py	1	# Copyright 2022 The JAX Authors.
LOW	tests/multi_device_test.py	1	# Copyright 2019 The JAX Authors.
LOW	tests/magma_linalg_test.py	1	# Copyright 2024 The JAX Authors.
LOW	tests/lax_numpy_setops_test.py	1	# Copyright 2025 The JAX Authors.
1339 more matches not shown…

Decorative Section Separators125 hits · 375 pts

Severity	File	Line	Snippet
MEDIUM	ci/run_pytest_tpu.sh	15	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	15	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	41	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	43	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	50	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	53	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	106	# ==============================================================================
MEDIUM	ci/run_pytest_rocm.sh	108	# ==============================================================================
MEDIUM	ci/run_bazel_test_cuda_non_rbe.sh	15	# ==============================================================================
MEDIUM	ci/run_bazel_cuda_targeted_tests.sh	15	# ==============================================================================
MEDIUM	ci/run_bazel_test_rocm_rbe.sh	15	# ==============================================================================
MEDIUM	ci/build_artifacts.sh	15	# ==============================================================================
MEDIUM	ci/build_rocm_artifacts.sh	15	# ==============================================================================
MEDIUM	ci/run_bazel_test_cpu_rbe.sh	15	# ==============================================================================
MEDIUM	ci/upload_rocm_logs.sh	15	# ==============================================================================
MEDIUM	ci/run_bazel_test_cuda_rbe.sh	15	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	15	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	45	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	47	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	54	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	57	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	107	# ==============================================================================
MEDIUM	ci/run_pytest_cuda.sh	109	# ==============================================================================
MEDIUM	ci/run_bazel_test_tpu.sh	15	# ==============================================================================
MEDIUM	ci/run_pytest_cpu.sh	15	# ==============================================================================
MEDIUM	ci/utilities/install_wheels_locally.sh	15	# ==============================================================================
MEDIUM	ci/utilities/run_docker_container.sh	15	# ==============================================================================
MEDIUM	ci/utilities/convert_msys_paths_to_win_paths.py	14	# ==============================================================================
MEDIUM	ci/utilities/collect_bazel_test_xmls.sh	15	# ==============================================================================
MEDIUM	ci/utilities/setup_build_environment.sh	15	# ==============================================================================
MEDIUM	ci/utilities/run_auditwheel.sh	15	# ==============================================================================
MEDIUM	ci/postprocess/process_test_results.sh	15	# ==============================================================================
MEDIUM	tests/mesh_utils_test.py	14	# ==============================================================================
MEDIUM	tests/hijax_test.py	2041	# ------------
MEDIUM	tests/hijax_test.py	2043	# ------------
MEDIUM	tests/hijax_test.py	2114	#------------
MEDIUM	tests/hijax_test.py	2116	#------------
MEDIUM	tests/mosaic/gpu_constraints_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_torch_distributed_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_dialect_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_torch_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/profiler_cupti_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/flash_attention_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_multidevice_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_layout_inference_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/matmul_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_test.py	14	# ==============================================================================
MEDIUM	tests/mosaic/gpu_distributed_test.py	14	# ==============================================================================
MEDIUM	tests/pallas/mgpu_ragged_dot_test.py	14	# ==============================================================================
MEDIUM	tests/pallas/mgpu_torch_test.py	14	# ==============================================================================
MEDIUM	tests/pallas/mgpu_examples_test.py	14	# ==============================================================================
MEDIUM	tests/pallas/mgpu_attention_test.py	14	# ==============================================================================
MEDIUM	tests/pallas/mgpu_collective_matmul_test.py	14	# ==============================================================================
MEDIUM	tests/pallas/mgpu_matmul_test.py	14	# ==============================================================================
MEDIUM	docs/notebooks/cute_dsl_jax/cute_dsl_jax_kernels.py	284	# ── Vector Add ────────────────────────────────────────────────────
MEDIUM	docs/notebooks/cute_dsl_jax/cute_dsl_jax_kernels.py	297	# ── SAXPY ─────────────────────────────────────────────────────────
MEDIUM	docs/notebooks/cute_dsl_jax/cute_dsl_jax_kernels.py	311	# ── ReLU ──────────────────────────────────────────────────────────
MEDIUM	docs/notebooks/cute_dsl_jax/cute_dsl_jax_kernels.py	323	# ── Fused Bias + ReLU ─────────────────────────────────────────────
MEDIUM	docs/notebooks/cute_dsl_jax/cute_dsl_jax_kernels.py	337	# ── GEMM ──────────────────────────────────────────────────────────
MEDIUM	docs/notebooks/cute_dsl_jax/cute_dsl_jax_kernels.py	351	# ── Elementwise Add (2-D) ─────────────────────────────────────────
65 more matches not shown…

Deep Nesting437 hits · 346 pts

Severity	File	Line
LOW	conftest.py	55
LOW	ci/postprocess/xml2json.py	70
LOW	tests/random_test.py	690
LOW	tests/export_back_compat_test.py	557
LOW	tests/lax_control_flow_test.py	2376
LOW	tests/errors_test.py	44
LOW	tests/ffi_test.py	45
LOW	tests/lax_numpy_einsum_test.py	295
LOW	tests/debug_info_test.py	119
LOW	tests/api_util_test.py	28
LOW	tests/dtypes_test.py	489
LOW	tests/fused_attention_stablehlo_test.py	123
LOW	tests/lax_numpy_test.py	2551
LOW	tests/lax_numpy_test.py	6579
LOW	tests/lax_numpy_reducers_test.py	164
LOW	tests/linalg_test.py	765
LOW	tests/linalg_test.py	1843
LOW	tests/linalg_test.py	1945
LOW	tests/linalg_test.py	1846
LOW	tests/linalg_test.py	1953
LOW	tests/profiler_test.py	54
LOW	tests/profiler_test.py	133
LOW	tests/lax_numpy_indexing_test.py	1491
LOW	tests/export_test.py	2391
LOW	tests/control_deps_test.py	53
LOW	tests/pgle_test.py	135
LOW	tests/pgle_test.py	221
LOW	tests/pgle_test.py	367
LOW	tests/sparse_bcoo_bcsr_test.py	57
LOW	tests/sparse_bcoo_bcsr_test.py	101
LOW	tests/device_test.py	24
LOW	tests/device_test.py	46
LOW	tests/compilation_cache_test.py	530
LOW	tests/transfer_guard_test.py	190
LOW	tests/transfer_guard_test.py	201
LOW	tests/api_test.py	4955
LOW	tests/api_test.py	5278
LOW	tests/api_test.py	5286
LOW	tests/documentation_coverage_test.py	140
LOW	tests/sparse_test.py	1021
LOW	tests/lax_test.py	117
LOW	tests/lax_test.py	1612
LOW	tests/lax_test.py	4565
LOW	tests/mosaic/gpu_test.py	662
LOW	tests/mosaic/gpu_test.py	1112
LOW	tests/mosaic/gpu_test.py	4993
LOW	tests/mosaic/gpu_test.py	5678
LOW	tests/mosaic/gpu_test.py	6862
LOW	tests/multiprocess/thread_guard_test.py	42
LOW	tests/multiprocess/pgle_test.py	44
LOW	tests/pallas/tpu_pallas_interpret_test.py	313
LOW	tests/pallas/tpu_pallas_interpret_test.py	409
LOW	tests/pallas/triton_pallas_test.py	142
LOW	…allas/tpu_fusible_matmul_with_stateful_fusions_test.py	84
LOW	tests/pallas/einshape_test.py	60
LOW	tests/pallas/ops_test.py	80
LOW	tests/pallas/ops_test.py	724
LOW	tests/pallas/ops_test.py	2126
LOW	tests/pallas/mosaic_gpu_test.py	3352
LOW	tests/pallas/tpu_fusible_matmul_test.py	84
377 more matches not shown…

Self-Referential Comments73 hits · 202 pts

Severity	File	Line	Snippet
MEDIUM	ci/utilities/run_docker_container.sh	55	# Create a temporary file to pass any user defined JAXCI_ / JAX_ / JAXLIB_
MEDIUM	ci/utilities/setup_build_environment.sh	105	# Create the output directory if it doesn't exist.
MEDIUM	tests/ragged_collective_test.py	769	# Define a mesh with PP + EP
MEDIUM	tests/lax_scipy_test.py	548	# This function is not defined for negative values, this makes sure they are nan
MEDIUM	tests/custom_api_test.py	2081	# Create the custom function
MEDIUM	tests/garbage_collection_guard_test.py	54	# Create a reference cycle of two jax.Arrays.
MEDIUM	tests/garbage_collection_guard_test.py	67	# Create a reference cycle of two jax.Arrays.
MEDIUM	tests/layout_test.py	361	# Create a custom layout instead of using `arr.layout` to test the API.
MEDIUM	tests/layout_test.py	384	# Create a custom layout instead of using `arr.layout` to test the API.
MEDIUM	tests/layout_test.py	665	# Create a custom layout instead of using `arr.layout` to test the API.
MEDIUM	tests/sparse_bcoo_bcsr_test.py	1153	# Create a matrix with duplicate indices
MEDIUM	tests/aot_test.py	337	# Create a compile-only topology, but DON'T switch to CPU so the real
MEDIUM	tests/array_test.py	666	# Create a few arrays
MEDIUM	tests/hijax_test.py	62	# Define a type
MEDIUM	tests/api_test.py	2053	# Creating an array from a numpy array with a fully-replicated sharding
MEDIUM	tests/api_test.py	2070	# Creating an array from a numpy array with a non-fully-replicated sharding
MEDIUM	tests/api_test.py	2086	# Creating an array from per-device JAX arrays calls internal
MEDIUM	tests/mosaic/gpu_layout_inference_test.py	1466	# Create a var to use in the constraint system.
MEDIUM	tests/multiprocess/array_test.py	613	# Create an array that is non-addressable in processes besides `pid`.
MEDIUM	tests/multiprocess/array_test.py	635	# Create a PRNG key array that is non-addressable in processes besides
MEDIUM	tests/multiprocess/array_test.py	669	# Create a sharding that is non-addressable in processes besides `pid`.
MEDIUM	tests/multiprocess/array_test.py	696	# Create a sharding that is non-addressable in processes besides `pid`.
MEDIUM	tests/multiprocess/array_test.py	723	# Create a single device sharding for a device local to process `pid`.
MEDIUM	tests/pallas/tpu_pallas_interpret_distributed_test.py	63	# Create an input array that shards the last dimension across
MEDIUM	tests/pallas/tpu_pallas_interpret_distributed_test.py	149	# Create an input array that shards the first dimension across
MEDIUM	tests/pallas/mosaic_gpu_test.py	6859	# Create an index-invariant output.
MEDIUM	tests/pallas/pallas_test.py	2509	# Create a validity mask for OOB values.
MEDIUM	docs/sphinxext/jax_list_config_options.py	38	# Create a field list item
MEDIUM	docs/sphinxext/jax_list_config_options.py	41	# Create the field name (label)
MEDIUM	docs/sphinxext/jax_list_config_options.py	46	# Create the field body (content)
MEDIUM	docs/sphinxext/jax_list_config_options.py	81	# Create a section for this option
MEDIUM	docs/sphinxext/jax_list_config_options.py	86	# Create a title with the option name (important for TOC)
MEDIUM	docs/sphinxext/jax_list_config_options.py	92	# Create a field list for side-by-side display
MEDIUM	jaxlib/tools/build_wheel.py	225	# This file is required by PEP-561. It marks jaxlib as package containing
MEDIUM	.github/workflows/bazel_cuda_b200_mosaic_presubmit.yml	87	# Create an array with the file patterns and
MEDIUM	jax/version.py	15	# This file is included as part of both jax and jaxlib. It is also
MEDIUM	…ops/tpu/splash_attention/splash_attention_mask_info.py	587	# Create a collection of the unique head masks in the input multi-head mask.
MEDIUM	…llas/ops/tpu/splash_attention/splash_attention_mask.py	374	# Define the mask function for chunk attention
MEDIUM	jax/experimental/pallas/ops/tpu/megablox/gmm.py	180	# Create the group ids for each grid index based on the tile counts for each
MEDIUM	jax/experimental/pallas/ops/tpu/megablox/gmm.py	230	# Create the m-dimension tile ids for each grid index based on the visit
MEDIUM	jax/experimental/pallas/ops/tpu/megablox/gmm.py	383	# Create the metadata we need for computation.
MEDIUM	jax/experimental/pallas/ops/tpu/megablox/gmm.py	632	# Create the metadata we need for computation.
MEDIUM	jax/experimental/jax2tf/call_tf.py	219	# Define the fwd and bwd custom_vjp functions
MEDIUM	jax/experimental/jax2tf/tests/jax2tf_test.py	349	x = tf.Variable(4.0, dtype=tf.float32) # Create a Tensorflow variable initialized to 4.0
MEDIUM	jax/experimental/jax2tf/examples/mnist_lib.py	225	# Create the model and save it
MEDIUM	jax/_src/test_util.py	2040	# The following function methods operate on mpmath number instances.
MEDIUM	jax/_src/custom_batching.py	178	# Define a class, instead of making a function closing over `rule`, so
MEDIUM	jax/_src/ad_checkpoint.py	400	# This function is similar to api_util.argnums_partial, except the error
MEDIUM	jax/_src/ad_checkpoint.py	926	>>> # Define a function where we explicitly name an intermediate value
MEDIUM	jax/_src/ffi.py	710	# This method is kept to support the behavior that was previously exposed
MEDIUM	jax/_src/interpreters/partial_eval.py	255	# Create the input tracers for the staged-out (unknown-value) call.
MEDIUM	jax/_src/interpreters/partial_eval.py	2449	# This function is conceptually the same thing as just calling eval_jaxpr,
MEDIUM	jax/_src/interpreters/mlir.py	1350	# Create a keepalives list that will be mutated during the lowering.
MEDIUM	jax/_src/lax/utils.py	15	# This module contains utility functions split out of jax._src.lax.lax to
MEDIUM	jax/_src/lax/control_flow/loops.py	1333	# Create the staged eqn.
MEDIUM	jax/_src/pallas/mosaic/lowering.py	3849	# Create a scalar constant.
MEDIUM	jax/_src/pallas/mosaic/lowering.py	3852	# Create a vector constant.
MEDIUM	jax/_src/pallas/mosaic_gpu/lowering.py	590	# The below code emission relies on the assumption that the first scratch
MEDIUM	jax/_src/lib/__init__.py	15	# This module is largely a wrapper around `jaxlib` that performs version
MEDIUM	jax/_src/export/serialization_generated.py	109	"""This method is deprecated. Please switch to GetRootAs."""
13 more matches not shown…

Cross-File Repetition27 hits · 135 pts

Severity	File	Snippet
HIGH	jaxlib/setup.py	this class makes 'bdist_wheel' include an abi tag on the wheel.
HIGH	jax_plugins/cuda/plugin_setup.py	this class makes 'bdist_wheel' include an abi tag on the wheel.
HIGH	jax_plugins/rocm/plugin_setup.py	this class makes 'bdist_wheel' include an abi tag on the wheel.
HIGH	jax_plugins/oneapi/plugin_setup.py	this class makes 'bdist_wheel' include an abi tag on the wheel.
HIGH	jax/experimental/mosaic/gpu/fragmented_array.py	returns the shape of the register array needed to represent an array of the given logical shape.
HIGH	jax/experimental/mosaic/gpu/fragmented_array.py	returns the shape of the register array needed to represent an array of the given logical shape.
HIGH	jax/experimental/mosaic/gpu/fragmented_array.py	returns the shape of the register array needed to represent an array of the given logical shape.
HIGH	…mental/jax2tf/tests/flax_models/transformer_nlp_seq.py	global hyperparameters used to minimize obnoxious kwarg plumbing.
HIGH	…erimental/jax2tf/tests/flax_models/transformer_lm1b.py	global hyperparameters used to minimize obnoxious kwarg plumbing.
HIGH	…perimental/jax2tf/tests/flax_models/transformer_wmt.py	global hyperparameters used to minimize obnoxious kwarg plumbing.
HIGH	jax/_src/export/shape_poly.py	returns -1 if self < other, 0 if self == other, 1 if self > other. the comparison is done lexicographically (syntactic),
HIGH	jax/_src/export/shape_poly.py	returns -1 if self < other, 0 if self == other, 1 if self > other. the comparison is done lexicographically (syntactic),
HIGH	jax/_src/export/shape_poly.py	returns -1 if self < other, 0 if self == other, 1 if self > other. the comparison is done lexicographically (syntactic),
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	jax/_src/export/serialization_generated.py	this method is deprecated. please switch to getrootas.
HIGH	…xport_back_compat_test_data/annotate_data_placement.py	#loc1 = loc("x") #loc2 = loc("y") module @jit_func attributes {jax.uses_shape_polymorphism = false, mhlo.num_partitions
HIGH	…xport_back_compat_test_data/annotate_data_placement.py	#loc1 = loc("x") #loc2 = loc("y") module @jit_func attributes {jax.uses_shape_polymorphism = false, mhlo.num_partitions
HIGH	…xport_back_compat_test_data/annotate_data_placement.py	#loc1 = loc("x") #loc2 = loc("y") module @jit_func attributes {jax.uses_shape_polymorphism = false, mhlo.num_partitions

Cross-Language Confusion24 hits · 130 pts

Severity	File	Line	Snippet
HIGH	tests/stack_test.py	34	stack = stack.push(jnp.int32(7))
HIGH	tests/stack_test.py	36	stack = stack.push(jnp.int32(8))
HIGH	tests/stack_test.py	41	stack = stack.push(jnp.int32(9))
HIGH	tests/export_test.py	1117	stack.push(self.assertRaisesRegex(Exception, re.escape(expect_error)))
HIGH	tests/export_test.py	934	stack.push(self.assertRaisesRegex(Exception, expect_error))
HIGH	tests/export_test.py	1040	stack.push(self.assertRaisesRegex(ValueError, expect_error_outer_exp))
HIGH	tests/export_test.py	1054	stack.push(self.assertRaisesRegex(Exception, expect_error_run))
HIGH	jax/experimental/jax2tf/tests/shape_poly_test.py	440	stack.push(self.assertRaisesRegex(Exception, re.escape(expect_error)))
HIGH	jax/experimental/colocated_python/func.py	363	func_backend.SINGLETON_RESULT_STORE.push(uid, result_leaves)
HIGH	jax/_src/jaxpr_util.py	627	while ((match = anchorRegex.exec(line)) !== null) {{
HIGH	jax/_src/jaxpr_util.py	672	const isCurrent = currentMatchIdx !== -1 && lineAbsoluteIdx === matchingLines[currentMatchIdx];
HIGH	jax/_src/jaxpr_util.py	705	let selectedElement = null;
HIGH	jax/_src/jaxpr_util.py	713	if (lineIdx !== undefined) {{
HIGH	jax/_src/jaxpr_util.py	738	while (currentIdx !== null && currentIdx !== undefined) {{
HIGH	jax/_src/jaxpr_util.py	738	while (currentIdx !== null && currentIdx !== undefined) {{
HIGH	jax/_src/jaxpr_util.py	743	renderedFrames.push({{file: file, func: func, line: frame.line, col: frame.col}});
HIGH	jax/_src/jaxpr_util.py	804	matchingLines.push(idx);
HIGH	jax/_src/jaxpr_util.py	808	matchingLines.push(idx);
HIGH	jax/_src/jaxpr_util.py	937	if (offsetRight > 100 && offsetRight < document.body.clientWidth - 100) {{
HIGH	jax/_src/lax/lax.py	9018	b <= a, then the result is undefined, and different implementations may
HIGH	jax/_src/numpy/linalg.py	529	discontinuous at det=0 so standard autodiff is undefined; we use sign_dot=0
HIGH	jax/_src/tpu/linalg/eigh.py	326	agenda = agenda.push(_Subproblem(offset=jnp.array(0, np.int32), size=n))
HIGH	jax/_src/tpu/linalg/eigh.py	417	agenda.push(_Subproblem(offset, rank)),
HIGH	jax/_src/tpu/linalg/eigh.py	436	agenda.push(_Subproblem(offset + rank, (b - rank))),

AI Slop Vocabulary47 hits · 114 pts

Severity	File	Line	Snippet
MEDIUM	ci/utilities/run_docker_container.sh	22	# run this script as they leverage built-in containerization features to run
MEDIUM	tests/export_back_compat_test.py	264	# Compute the inputs to simplify the harness
MEDIUM	tests/lax_numpy_operators_test.py	294	# probably add a custom test harness for unwrap that tests the period
LOW	tests/custom_api_test.py	366	# useful either: instead of using nondiff_argnums here, a user can just pass
MEDIUM	tests/lax_numpy_test.py	3350	# permissive canonicalization logic in the test harness.
MEDIUM	tests/export_harnesses_multi_platform_test.py	60	# If you want to run this test for only one harness, add parameter
MEDIUM	tests/shape_poly_test.py	3671	# Pick the dtype with the most harnesses in this group. Some harness
MEDIUM	tests/shape_poly_test.py	3743	# If you want to run this test for only one harness that includes "foo"
MEDIUM	tests/shape_poly_test.py	3800	# Update this here rather than in harness object because vmap_random_gamma is derived
LOW	tests/shape_poly_test.py	1862	# It is not sufficient to just use the shape of an input; it is still unused
MEDIUM	tests/linalg_test.py	984	# This expresses identity function, which makes us robust to, e.g., the
MEDIUM	tests/lax_autodiff_test.py	159	# TODO(mattjj): make some-equal checks more robust, enable second-order
MEDIUM	tests/sparsify_test.py	341	# Note: more comprehensive tests in sparse_test.py:test_bcoo_squeeze
MEDIUM	tests/sparsify_test.py	354	# Note: more comprehensive tests in sparse_test.py:test_bcoo_rev
MEDIUM	tests/sparsify_test.py	419	# Note: more comprehensive tests in sparse_test.py:test_bcoo_conv_general_dilated
LOW	tests/api_test.py	454	# We can't just use `lambda x: x` because JAX simplifies this away to an
LOW	tests/mosaic/gpu_layout_inference_test.py	90	# TMEM reference, so we can just return a trivial mapping.
LOW	docs/array_refs.py	115	# `jax.ref.swap`, but usually you'd just use NumPy-style array indexing syntax:
LOW	docs/autodidax.py	1475	# call `bind`, in the primitive wrapper we can just use `make_jaxpr` to form
MEDIUM	…ops/tpu/splash_attention/splash_attention_mask_info.py	892	# maintain the SPMD paradigm.
MEDIUM	…erimental/jax2tf/tests/jax_primitives_coverage_test.py	50	# If you want to run this test for only one harness, add parameter
MEDIUM	…erimental/jax2tf/tests/jax_primitives_coverage_test.py	90	# f"{[u.description for u in jax_unimpl]} in harness: {harness.fullname}"))
MEDIUM	jax/experimental/jax2tf/tests/shape_poly_test.py	223	# Makes and tests a harness. See PolyHarness documentation.
LOW	jax/experimental/jax2tf/tests/shape_poly_test.py	587	# It is not sufficient to just use the shape of an input; it is still unused
MEDIUM	jax/experimental/jax2tf/tests/model_harness.py	77	"""Partially apply harness in order to create variables lazily.
MEDIUM	jax/experimental/jax2tf/tests/tf_test_util.py	221	# Run JAX. Should not fail, we assume that the harness has been filtered
MEDIUM	jax/experimental/jax2tf/tests/primitives_test.py	91	# If you want to run this test for only one harness, add parameter
LOW	jax/_src/compiler.py	494	# simply return the non-PGLE profiled module from the persistent cache if it
LOW	jax/_src/pjit.py	1019	# checks and just return the pjit_in_shardings directly. `shard_args` will
LOW	jax/_src/checkify.py	348	# Default non-HOP case: just call primitive and don't update error.
LOW	jax/_src/interpreters/batching.py	632	# if there's only agreeing batch dims and scalars, just call the primitive
LOW	jax/_src/lax/ann.py	367	# 1. ApproxTopK is internally a variadic reduce, so we can simply call
MEDIUM	jax/_src/lax/control_flow/loops.py	2023	# lowering. Fundamentally, we'd like to rewrite a while loop that looks like
MEDIUM	jax/_src/numpy/linalg.py	1445	# TODO: add custom jvp rule for more robust lstsq differentiation
LOW	jax/_src/pallas/pallas_call.py	689	# dimensions. For now, we just use 0.
LOW	jax/_src/pallas/mosaic/lowering.py	1682	# `jax.empty_ref`), but lowering expects them to exist---so we just return
MEDIUM	…/_src/pallas/mosaic/interpret/interpret_pallas_call.py	1592	# TODO(nrink): It would be more robust if the buffer id, i.e. `src`,
LOW	…/_src/pallas/mosaic/interpret/interpret_pallas_call.py	636	# callback to `get`. Should we just pass the shape to `get`?
LOW	jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py	662	# callback to `get`. Should we just pass the shape to `get`?
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	148	# Descriptive name of the harness, used as a testcase_name. Unique in a group.
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	157	# partially implemented in JAX for this harness.
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	386	# Change the testcase name to include the harness name.
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	2003	# This first harness runs the tests for all dtypes using default values for
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	2589	# This first harness runs the tests for all dtypes using default values for
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	2894	# This first harness runs the tests for all dtypes and precisions using
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	3095	# This first harness runs the tests for all dtypes and precisions using
MEDIUM	jax/_src/internal_test_util/test_harnesses.py	351	self.enabled = enabled # Does it apply to the current harness?

Dead Code55 hits · 104 pts

Severity	File	Line
MEDIUM	tests/custom_api_test.py	374
MEDIUM	tests/custom_api_test.py	376
MEDIUM	tests/custom_api_test.py	379
MEDIUM	tests/custom_api_test.py	382
MEDIUM	tests/custom_api_test.py	385
MEDIUM	tests/custom_api_test.py	386
MEDIUM	tests/custom_api_test.py	387
MEDIUM	tests/custom_api_test.py	1816
MEDIUM	tests/custom_api_test.py	1828
MEDIUM	tests/custom_api_test.py	1831
MEDIUM	tests/custom_api_test.py	1832
MEDIUM	tests/custom_api_test.py	1833
MEDIUM	tests/custom_api_test.py	1835
MEDIUM	tests/custom_api_test.py	1836
MEDIUM	tests/custom_api_test.py	1837
MEDIUM	tests/lax_numpy_reducers_test.py	907
MEDIUM	tests/lax_numpy_reducers_test.py	908
MEDIUM	tests/lax_numpy_reducers_test.py	909
MEDIUM	tests/lax_numpy_reducers_test.py	919
MEDIUM	tests/lax_numpy_reducers_test.py	920
MEDIUM	tests/lax_numpy_reducers_test.py	921
MEDIUM	tests/pmap_test.py	2067
MEDIUM	tests/pmap_test.py	2069
MEDIUM	tests/pmap_test.py	2075
MEDIUM	tests/pmap_test.py	2085
MEDIUM	tests/pmap_test.py	2089
MEDIUM	tests/pmap_test.py	2104
MEDIUM	tests/pmap_test.py	2110
MEDIUM	tests/pmap_test.py	2111
MEDIUM	tests/pmap_test.py	2112
MEDIUM	tests/pmap_test.py	2113
MEDIUM	tests/pmap_test.py	2124
MEDIUM	tests/pmap_test.py	2130
MEDIUM	tests/pmap_test.py	2131
MEDIUM	tests/pmap_test.py	2132
MEDIUM	tests/pmap_test.py	2133
MEDIUM	tests/pmap_test.py	2139
MEDIUM	tests/pmap_test.py	2145
MEDIUM	tests/pmap_test.py	2295
MEDIUM	tests/pmap_test.py	2297
MEDIUM	tests/pmap_test.py	2300
MEDIUM	tests/pmap_test.py	2301
MEDIUM	tests/pmap_test.py	2303
MEDIUM	tests/pmap_test.py	2304
MEDIUM	tests/pmap_test.py	2305
MEDIUM	tests/pmap_test.py	2306
MEDIUM	tests/pmap_test.py	2318
MEDIUM	tests/lax_test.py	1774
MEDIUM	tests/lax_test.py	1806
MEDIUM	tests/mosaic/gpu_test.py	949
MEDIUM	tests/mosaic/gpu_test.py	951
MEDIUM	tests/mosaic/gpu_test.py	952
MEDIUM	tests/mosaic/gpu_test.py	953
MEDIUM	tests/mosaic/gpu_test.py	960
MEDIUM	jax/_src/pallas/core.py	1675

Excessive Try-Catch Wrapping86 hits · 95 pts

Severity	File	Line	Snippet
MEDIUM	ci/utilities/convert_msys_paths_to_win_paths.py	40	print("Error: cygpath not found. Make sure it's in your PATH.")
MEDIUM	ci/utilities/convert_msys_paths_to_win_paths.py	43	print(f"Error converting path: {e}")
LOW	ci/postprocess/xml2json.py	87	except Exception as e:
LOW	tests/array_interoperability_test.py	155	except Exception as e:
LOW	tests/debug_info_test.py	111	except Exception as e:
LOW	tests/shape_poly_test.py	2041	except Exception:
LOW	tests/profiler_test.py	63	except Exception as e:
LOW	tests/transfer_guard_test.py	117	except Exception as e:
LOW	tests/transfer_guard_test.py	136	except Exception as e:
LOW	tests/api_test.py	5039	except Exception as e:
LOW	tests/pallas/gpu_pallas_distributed_test.py	600	except Exception:
LOW	tests/pallas/gpu_pallas_distributed_test.py	651	except Exception:
LOW	tests/pallas/gpu_pallas_distributed_test.py	1306	except Exception:
LOW	tests/pallas/ops_test.py	709	except Exception as e:
LOW	tests/pallas/ops_test.py	837	except Exception as e:
LOW	tests/pallas/ops_test.py	2688	except Exception as e:
LOW	docs/array_refs.py	82	except Exception as e:
LOW	docs/array_refs.py	222	except Exception as e:
LOW	docs/array_refs.py	231	except Exception as e:
LOW	docs/array_refs.py	240	except Exception as e:
LOW	docs/array_refs.py	249	except Exception as e:
LOW	docs/array_refs.py	289	except Exception as e:
LOW	docs/parallel.md	425	except Exception as e:
LOW	docs/parallel.md	464	except Exception as e:
LOW	docs/array_refs.md	87	except Exception as e:
LOW	docs/array_refs.md	222	except Exception as e:
LOW	docs/array_refs.md	231	except Exception as e:
LOW	docs/array_refs.md	240	except Exception as e:
LOW	docs/array_refs.md	249	except Exception as e:
LOW	docs/array_refs.md	288	except Exception as e:
LOW	docs/parallel.py	379	except Exception as e:
LOW	docs/parallel.py	417	except Exception as e:
LOW	docs/sphinxext/source_include.py	72	except Exception as e:
MEDIUM	docs/sphinxext/source_include.py	40	def get_tagged_block(filepath, tag, lines_spec=None):
LOW	…atic/fault_tolerance/data_parallelism_with_recovery.py	171	except Exception as e:
LOW	docs/_static/fault_tolerance/live_devices.py	56	except Exception as e:
LOW	docs/_static/fault_tolerance/data_parallelism.py	118	except Exception as e:
MEDIUM	docs/pallas/tpu/sparse.md	148	print("Error \|result - lax.dynamic_slice\| =", diff)
LOW	jax/__init__.py	28	except Exception as exc:
LOW	jax/experimental/multihost_utils.py	679	except Exception as e:
LOW	jax/experimental/array_serialization/serialization.py	238	except Exception as e:
MEDIUM	jax/experimental/array_serialization/serialization.py	197	def _thread_func(self):
LOW	jax/experimental/jax2tf/call_tf.py	294	except Exception as e:
LOW	jax/experimental/jax2tf/call_tf.py	591	except Exception as e:
LOW	…erimental/jax2tf/tests/jax_primitives_coverage_test.py	72	except Exception as e:
LOW	jax/experimental/jax2tf/tests/models_test_main.py	206	except Exception as e:
LOW	jax/experimental/jax2tf/tests/tf_test_util.py	246	except Exception as e:
LOW	jax/experimental/jax2tf/tests/primitives_test.py	135	except Exception as e:
LOW	jax/experimental/jax2tf/examples/mnist_lib.py	76	except Exception as e:
LOW	jax/experimental/colocated_python/obj_backend.py	107	except Exception as exc:
LOW	jax/experimental/sparse/bcsr.py	892	except Exception:
LOW	jax/_src/environment_info.py	30	except Exception:
MEDIUM	jax/_src/environment_info.py	27	def try_nvidia_smi() -> str \| None:
LOW	jax/_src/xla_bridge.py	832	except Exception as err:
LOW	jax/_src/compiler.py	83	except Exception:
LOW	jax/_src/compiler.py	760	except Exception as ex:
LOW	jax/_src/compiler.py	779	except Exception as ex:
LOW	jax/_src/compiler.py	830	except Exception as ex:
LOW	jax/_src/util.py	540	except Exception:
LOW	jax/_src/util.py	635	except Exception as e:
26 more matches not shown…

Redundant / Tautological Comments26 hits · 38 pts

Severity	File	Line	Snippet
LOW	tests/array_interoperability_test.py	183	# Check if the source device is preserved
LOW	tests/nn_test.py	119	# Check if float8_e8m0fnu is available
LOW	tests/nn_test.py	327	# Check if cudnn backend is called (only on CUDA).
LOW	tests/pgle_test.py	277	# Check if FDO profile file of the biggest module is not empty
LOW	tests/pgle_test.py	285	# Check if FDO profile file in dump directory is not empty
LOW	tests/multiprocess/array_test.py	297	# Check if we can specify that local input actually contains full-span
LOW	…allas/tpu_fusible_matmul_with_stateful_fusions_test.py	145	# Check if either input fusion reads from a Ref that is also written by the
LOW	tests/pallas/tpu_pallas_call_print_test.py	149	# Check if the numbers in the output match the values generated by `arange`.
LOW	.github/workflows/upload_metadata.yml	66	# Check if zip file is empty before unzipping
LOW	.github/workflows/verify-squash.yml	31	# Check if the skip label is present
LOW	.github/workflows/verify-squash.yml	54	# Check if squashed
LOW	.github/workflows/build_artifacts.yml	171	# Set shell to cmd to avoid path errors when using gcloud commands on Windows
LOW	jax/experimental/mosaic/gpu/fragmented_array.py	4843	# Check if input maps exactly to the end (prevents trailing dims).
LOW	jax/_src/test_util.py	1276	# Check if strict hermeticity is already satisfied
LOW	jax/_src/checkify.py	865	# Check if the first cond application will error.
LOW	jax/_src/pretty_printer.py	48	# Check if we're in IPython or Colab
LOW	jax/_src/lax/control_flow/loops.py	2200	# Check if the same Ref is written to in both cond and body.
LOW	jax/_src/state/primitives.py	763	# Check if start is static (which it can be)
LOW	jax/_src/state/primitives.py	776	# Check if we are indexing with a scalar or not. If we are indexing
LOW	jax/_src/numpy/setops.py	177	# Set mask to zero at locations corresponding to unique() padding.
LOW	jax/_src/numpy/array_constructors.py	251	# Check if object supports any of the data exchange protocols
LOW	jax/_src/pallas/core.py	1231	elif grid and isinstance(grid[0], tuple): # Check if we have a named grid
LOW	jax/_src/pallas/einshape.py	230	# Check if all RHS dims are known
LOW	jax/_src/scipy/signal.py	681	# Check if we can broadcast the outer axes together
LOW	jax/_src/scipy/signal.py	730	# Check if x and y are the same length, zero-pad if necessary
LOW	jax/_src/cudnn/fused_attention_stablehlo.py	1260	# Check if all required keys are present

Hallucination Indicators3 hits · 30 pts

Severity	File	Line	Snippet
CRITICAL	tests/mosaic/profiler_cupti_test.py	96	jax._src.lib.mosaic_gpu._mosaic_gpu_ext._cupti_init()
CRITICAL	tests/filecheck/jax_mlir_ext.filecheck.py	114	return jax._src.lib._jax.Traceback.get_traceback()
CRITICAL	jax/_src/pallas/mosaic_gpu/lowering.py	3910	for k, v in ctx.launch_ctx.module.operation.attributes.items():

Synthetic Comment Markers3 hits · 20 pts

Severity	File	Line	Snippet
HIGH	docs/contributing.md	26	## Can I contribute AI generated code?
HIGH	docs/contributing.md	35	You are responsible for any code you contribute to JAX, regardless of whether it was written manually or generated by AI
HIGH	jaxlib/py_executable.h	287	// Python objects to keep alive as requested by user.

Slop Phrases9 hits · 14 pts

Severity	File	Line	Snippet
MEDIUM	docs/parallel.py	396	# want the compiler to choose them automatically, you can use the `@auto_axes`
MEDIUM	jax/experimental/jax2tf/examples/saved_model_lib.py	102	# names, you can use `tree.map_structure_with_path` from the `dm-tree` package
LOW	jax/_src/core.py	2821	# but we make sure to reset it to Device because the Ref owns the memory space
MEDIUM	jax/_src/ad_checkpoint.py	272	Here is a simple example:
MEDIUM	jax/_src/errors.py	201	indices, such as with :code:`.at[...].set(...)`. Here is a simple example::
MEDIUM	jax/_src/errors.py	666	Here is a simple example of code that would lead to such an error::
MEDIUM	jax/_src/lax/windowed_reductions.py	154	Here is a simple example of a windowed product over pairs in a 1-dimensional array:
MEDIUM	jax/_src/lax/slicing.py	373	For example, here is how you can extract values at particular indices using
LOW	jax/_src/pallas/core.py	782	# If you hit this, make sure you take transforms into account and use either

Overly Generic Function Names9 hits · 10 pts

Severity	File	Line	Snippet
LOW	tests/core_test.py	584	def my_function():
LOW	tests/util_test.py	80	def my_function():
LOW	tests/aot_test.py	122	def my_function(x):
LOW	tests/aot_test.py	143	def my_function(x):
LOW	tests/api_test.py	93	def my_function():
LOW	tests/api_test.py	107	def my_function():
LOW	tests/api_test.py	1632	def my_function(x, flag):
LOW	jax/typing.py	47	def my_function(x: ArrayLike) -> Array:
LOW	jax/experimental/jax2tf/tests/jax2tf_test.py	306	def test_function(self):

Example Usage Blocks3 hits · 5 pts

Severity	File	Line	Snippet
LOW	ci/utilities/run_docker_container.sh	24	# Usage:
LOW	jaxlib/mosaic/dialect/tpu/util.h	110	// Example usage:
LOW	jax/_src/deprecations.py	22	# Example usage:

Verbosity Indicators1 hit · 2 pts

Severity	File	Line	Snippet
LOW	jaxlib/python_ref_manager.h	85	// The purpose of this function is to amortize lock acquisition costs over

Score History

Severity Breakdown

Pattern Findings