[TEST] Migrate AMSS to JAX

[HumphreyYang]

The aim of this issue is to migrate the amss lecture to use JAX.

Specifically:

1. Follow the guidelines listed at https://manual.quantecon.org/styleguide/jax-conversion.html.
2. Break up large classes and use autodifferentiation for marginal utility calculations. Some of the classes are not currently included in the lecture (e.g., AMSS classes and utility classes). These can be found in lectures/_static/lecture_specific/amss. Please add them to the lecture using a proper NamedTuple structure, following the guidelines.
3. Avoid using methods within classes. Move them out as pure functions that take parameter classes as arguments.
4. Correct any typos and mathematical errors in the lecture.
5. Follow the existing lecture setup, keep the implementation concise, and ensure the code closely matches the mathematical notation. Optimize the code using vectorization and JIT whenever possible.

[HumphreyYang]

Hi @mmcky, I opened this test issue as we discussed! Let's see how copilot performs.

It would probably give us a result that is 50-60% complete.

[mmcky]

@HumphreyYang -- nicely written issue!