Test of rates

[shangyangli]

No description provided.

[chaoming0625]

The tests have failed. Before merging this PR, you should fix the errors first.

Particularly, the bug fixing should be:

[chaoming0625]

At the line of

if bm.any(self.noise_i != 0.):

change it as:

with jax.ensure_compile_time_eval():
  has_noise = bm.any(self.noise_i != 0.)

if has_noise:
  ...

Similarly, at the line of ```` if bm.any(self.noise_e != 0.):``, change it as:

with jax.ensure_compile_time_eval():
  has_noise = bm.any(self.noise_e != 0.)

if has_noise:
  ...

[chaoming0625]

@ztqakita Please help @shangyangli to pass all tests. We cannot merge any PR with test errors.

[ztqakita]

Great!