Fix runtime crash with the local_subtensor_of_alloc optimization when the input have a broadcastable dimensions.