nn: use ensure_arraylike instead of check_arraylike and jnp.asarray

[arvindsankar]

From poking around, it seems like invoking jnp.asarray after check_arraylike duplicates the checking work and could be expensive in some cases [1]. From a previous PR [2], I see there seems to be a primitive that checks an input is arraylike and returning it back as a jnp array.

[1] #25931
[2] info gained from @jakevdp 's comments on #32668.

This might be a tiny performance win since this feels like a hot path to me. I'm happy to benchmark, but I'm not sure how to or if thats needed.

[gemini-code-assist]

Summary of Changes

Hello @arvindsankar, I'm Gemini Code Assist¹! 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 streamlines array input handling within JAX's neural network functions. By consolidating the check_arraylike validation and jnp.asarray conversion into a single ensure_arraylike utility, the change aims to reduce redundant operations and potentially enhance the performance of these core activation functions. This improves code efficiency and consistency across the affected functions.

Highlights

• Code Refactoring: Replaced the two-step process of numpy_util.check_arraylike followed by jnp.asarray with a single call to numpy_util.ensure_arraylike across multiple neural network activation functions.
• Performance Optimization: The change aims to eliminate redundant array checking and conversion, potentially leading to a minor performance improvement, especially in frequently called 'hot path' functions.
• Consistency: Standardized array handling within jax/_src/nn/functions.py by adopting a more efficient and consolidated utility function.

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. ↩

[gemini-code-assist]

Code Review

This pull request correctly replaces separate calls to check_arraylike and jnp.asarray with a single ensure_arraylike call across various functions in jax/_src/nn/functions.py. This change simplifies the code and should provide a minor performance improvement by avoiding redundant checks. The changes are consistent and well-executed. I have one suggestion for the squareplus function to further improve efficiency by handling both of its array-like arguments in a single function call.

[jakevdp]

Thanks!