Add caching to trace_to_subjaxpr_dynamic.

[pschuh]

This allows the MLIR lowering code to cache call lowerings.

example output:

module @jit_fun.0 {
  func.func public @main(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
    %0 = call @square(%arg0) : (tensor<4x8xf32>) -> tensor<4x8xf32>
    %1 = call @square(%0) : (tensor<4x8xf32>) -> tensor<4x8xf32>
    return %1 : tensor<4x8xf32>
  }
  func.func private @square(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
    %0 = mhlo.multiply %arg0, %arg0 : tensor<4x8xf32>
    return %0 : tensor<4x8xf32>
  }
}

If / when jaxprs support recursion, this approach will still work because the mlir lowering cache operates on Jaxpr object identity.

[mattjj]

Thank you!