Skip to content

Comments

Another batch of Pyrefly fixes#35332

Merged
copybara-service[bot] merged 1 commit intojax-ml:mainfrom
superbobry:pyrefly
Feb 23, 2026
Merged

Another batch of Pyrefly fixes#35332
copybara-service[bot] merged 1 commit intojax-ml:mainfrom
superbobry:pyrefly

Conversation

@superbobry
Copy link
Collaborator

No description provided.

@superbobry superbobry requested a review from jakevdp February 23, 2026 14:15
@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @superbobry, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request primarily focuses on enhancing the type-checking compliance of the JAX codebase, particularly in relation to the pyrefly static analysis tool. The changes involve a comprehensive set of type annotation adjustments, the strategic placement of pyrefly: ignore comments to address specific type inference challenges, and minor code modifications to align with stricter type enforcement. The overall goal is to improve code maintainability and robustness by ensuring type correctness without altering the existing runtime behavior or core functionality.

Highlights

  • Type Annotation Refinements: Updated various type hints across the codebase, including changing Sequence to Iterable, adding TypeAlias for clarity, and providing explicit None type unions for optional variables.
  • Pyrefly Type-Checking Ignored: Introduced numerous pyrefly: ignore comments to suppress specific type-checking errors, addressing issues like missing-attribute, not-iterable, bad-argument-type, unsupported-operation, and missing-import.
  • Positional-Only Arguments: Added positional-only argument markers (/) to __call__ methods within _SerializeAuxData and _DeserializeAuxData protocols to enforce argument passing style.
  • Code Clarity and Assertions: Added assert statements and del record for unused variables to improve code clarity and satisfy static analysis tools.
  • Minor Logic Adjustments: Made small functional changes such as replacing np.int32(0.) with np.zeros((), dtype=dtype) for better type consistency in test harnesses and refactoring context management in _lower_mosaic_module_to_asm.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • jax/_src/export/_export.py
    • Added positional-only argument markers to __call__ methods in _SerializeAuxData and _DeserializeAuxData protocols.
    • Added pyrefly: ignore[missing-attribute] to json.dumps(nodetype._fields) and hlo.get_version_from_compatibility_requirement calls.
    • Added pyrefly: ignore[not-iterable] to a loop over arg_attrs[idx].
    • Added pyrefly: ignore[bad-argument-type] to _get_named_sharding call.
    • Added pyrefly: ignore[missing-attribute] to sdy.MeshAttr and sdy.MeshAxisAttr calls.
  • jax/_src/export/serialization.py
    • Replaced Sequence with Iterable in the _serialize_array function signature.
    • Imported cast and used it for type assertion on in_shardings_hlo and out_shardings_hlo.
    • Added an assert node_type is not None statement in _serialize_pytree_node.
    • Added pyrefly: ignore[bad-argument-type] to a call to _serialize_array.
    • Refactored _deserialize_aval to directly return the result of core.update_aval_with_sharding.
  • jax/_src/export/serialization_generated.py
    • Added a pyrefly: ignore-errors comment at the file level.
  • jax/_src/export/shape_poly.py
    • Imported TypeAlias from typing.
    • Used TypeAlias for SortedTerms and SortedFactors.
    • Renamed parameters e1 and e2 to pairs1 and pairs2 in _linear_combination_sorted_pairs overloads.
    • Changed the return type annotation of _divmod from tuple[DimSize, int] to tuple[DimSize, DimSize].
    • Added explicit type hints DimSize | None for acc variables in expr and term methods.
  • jax/_src/internal_test_util/export_back_compat_test_data/cpu_triangular_solve_blas_trsm.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/cpu_tridiagonal_solve_lapack_gtsv.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/cuda_cholesky_solver_potrf.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/cuda_eigh_cusolver_syev.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/cuda_lu_cusolver_getrf.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/cuda_qr_cusolver_geqrf.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/cuda_tridiagonal_cusolver_sytrd.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/rocm_cholesky_solver_potrf.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/rocm_lu_rocsolver_getrf.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/rocm_qr_hipsolver_geqrf.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/export_back_compat_test_data/rocm_tridiagonal_hipsolver_sytrd.py
    • Added a pyrefly: ignore-errors comment.
  • jax/_src/internal_test_util/test_harnesses.py
    • Added a type hint set[str] to the names variable.
    • Added pyrefly: ignore[unsupported-operation] to set difference operations involving jtu.dtypes.all.
    • Changed the initialization of an element in init_values from np.int32(0.) to np.zeros((), dtype=dtype).
  • jax/_src/lax/convolution.py
    • Added pyrefly: ignore[missing-attribute] to hlo.ConvDimensionNumbers.get.
  • jax/_src/lax/fft.py
    • Added pyrefly: ignore[missing-attribute] to hlo.FftOp.
  • jax/_src/lazy_loader.py
    • Added an assert owner_name is not None statement with a pyrefly#40 comment.
  • jax/_src/lib/init.py
    • Added pyrefly: ignore[missing-import] to the from mlir import ir import statement within an except block.
  • jax/_src/logging_config.py
    • Added del record and a type hint for the record parameter in _DebugHandlerFilter.filter.
  • jax/_src/mesh_utils.py
    • Added a type hint np.ndarray | None to the best_logical_axis_assignment variable.
  • jax/_src/pallas/mosaic_gpu/lowering.py
    • Added an assert isinstance(ctx.out_layout_hint, mgpu.TiledLayout) statement.
  • jax/_src/path.py
    • Updated the return type annotation of make_jax_dump_dir from pathlib.Path | None to Path | None.
  • jax/_src/test_util.py
    • Added pyrefly: ignore[missing-attribute] to hlo.get_smaller_version.
  • jax/_src/tpu_custom_call.py
    • Modified result_types assignment to use mlir.flatten_ir_types with map(mlir.aval_to_ir_type, out_avals).
    • Added a type hint dict[str, ir.Attribute] | None to the extra_attributes variable.
    • Changed the return type annotation of _lower_mosaic_module_to_asm from tuple[ir.Module, ...] to tuple[bytes, ...].
    • Refactored the context management for module.context in _lower_mosaic_module_to_asm.
  • jax/_src/traceback_util.py
    • Added pyrefly: ignore[unsupported-operation] to a comparison involving IPython.version_info.
Activity
  • The pull request was authored by superbobry.
  • The changes primarily involve adding pyrefly: ignore comments and refining type annotations across numerous files to improve static analysis compliance.
  • Minor code adjustments were made to align with type-checking requirements or improve clarity, such as adding assert statements and updating variable initializations.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a batch of fixes for the Pyrefly static analysis tool. Most changes involve adding # pyrefly: ignore comments to suppress warnings or improving type hints, which are reasonable. However, I've identified one change in jax/_src/internal_test_util/test_harnesses.py that appears to introduce a bug by using an incorrect dtype for an initial value in a lax.reduce call within a test harness. I've provided a specific comment and a suggested fix for this issue.

init_values = [init_val]
if nr_operands == 2:
init_values.append(np.int32(0.))
init_values.append(np.zeros((), dtype=dtype))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The initial value for the second operand in this lax.reduce call should have a dtype of np.int32, as the second operand is always an int32 array in this test harness. The current implementation uses dtype, which is the dtype of the first operand and can be a float. This will cause a type mismatch when dtype is not an integer type.

The previous implementation np.int32(0.) was correct. To fix this, you should use a hardcoded np.int32 dtype for the second initial value.

Suggested change
init_values.append(np.zeros((), dtype=dtype))
init_values.append(np.zeros((), dtype=np.int32))

elements: Sequence[T],
elements: Iterable[T],
) -> int:
element_offsets = [serialize_one(builder, e) for e in elements]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since elements is now an interable rather than a sequence, I'd consider adding del elements in the next line to emphasize that it may be consumed by this iteration.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good idea, done.

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Feb 23, 2026
@superbobry superbobry force-pushed the pyrefly branch 2 times, most recently from 37b4965 to 3aa1911 Compare February 23, 2026 15:37
@copybara-service copybara-service bot merged commit b486aac into jax-ml:main Feb 23, 2026
25 of 31 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kokoro:force-run pull ready Ready for copybara import and testing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants