Skip to content

Move dtypes.canonicalize_value to C++.#38171

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_926256875
Open

Move dtypes.canonicalize_value to C++.#38171
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_926256875

Conversation

@copybara-service
Copy link
Copy Markdown

Move dtypes.canonicalize_value to C++.

Since this function is primarily "apply table of handlers", we can speed it up both by moving to C++ both:

  • the handler finding code, and
  • some of the most common handlers, notably the identity function.

Notes:

  • We retain the ability to register handlers in Python for novel types.
  • We now want to call a weak key weak value cache from C++ when canonicalizing np.ndarray objects, so we add a C++ API for that also.
  • Since calling canonicalize_dtype is now faster than skipping it with a bunch of isinstance checks, we add a number of additional types (Tracer, ShapeDtypeStruct) to the set of things that canonicalize_value knowns how to handle and call it unconditionally in places like Primitive.bind().

Since this function is primarily "apply table of handlers", we can speed it up both by moving to C++ both:
* the handler finding code, and
* some of the most common handlers, notably the identity function.

Notes:
* We retain the ability to register handlers in Python for novel types.
* We now want to call a weak key weak value cache from C++ when canonicalizing np.ndarray objects, so we add a C++ API for that also.
* Since calling canonicalize_dtype is now faster than skipping it with a bunch of isinstance checks, we add a number of additional types (Tracer, ShapeDtypeStruct) to the set of things that canonicalize_value knowns how to handle and call it unconditionally in places like Primitive.bind().
PiperOrigin-RevId: 926256875
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant