M4ML C2 Week 1 Assignment 1 exercise 3 problem

I’m having a type error issue with the f_of_omege function when it’s being called by L_of_omega_array function in exercise 3. I’m getting the error:

TypeError: unsupported operand type(s) for *: 'DeviceArray' and '_IndexUpdateRef'

it’s referencing “prices_A * omega”
I think it has something to do with switching to JAX but am having trouble fixing it.

So, for others’ reference out there. the .at() in the first line of code in exercise 3 does not work. When I did a normal index of the omega array it worked.