Fix deprecated usage of JAX symbols.

[ethanluoyc]

Description

Remove deprecated JAX symbols.

Motivation and Context

• jax.abstract_array.ShapedArray is removed for jax>=0.4.16. See

https://github.com/google/jax/blob/main/CHANGELOG.md#jax-0416-sept-18-2023

• xla_client.register_cpu_custom_call_target is removed in

openxla/xla@0ab5486

close #281

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Put an x in all the boxes that apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• New environment (non-breaking change which adds 3rd-party environment)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of example)

Implemented Tasks

• Fix deprecated symbol

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code using make format (required)
• I have checked the code using make lint (required)
• I have ensured make bazel-test pass. (required)

[mavenlin]

@ethanluoyc Thanks for fixing the bug!
The code change looks good to me, but it seems some tests are broken.

[Trinkle23897]

I’ll fix it tmr