Indexing Error: InvalidArgumentError: slice index 1000 of dimension 0 out of bounds. [Op:StridedSlice] name: fit_sparse/minimize/while/minimize_one_step/while/strided_slice/

Receiving the below error when trying to run fit_sparse on GLM model where feature dimension is larger than sample dimension:

InvalidArgumentError: slice index 1000 of dimension 0 out of bounds. [Op:StridedSlice] name: fit_sparse/minimize/while/minimize_one_step/while/strided_slice/

Minimum example:

import tensorflow_probability as tfp
import tensorflow as tf

print(tf.__version__) #== 2.0.0

print(tfp.__version__) #== 0.8.0


n_samples = 1000
n_features = 2000

model = tfp.glm.Bernoulli()
x_ = tf.Variable(np.random.normal(0,1,(n_samples,n_features)), dtype=np.float32)
y_ = tf.Variable(np.random.binomial(1,.3,(n_samples,)), dtype=np.float32)
model_coefficients_start = tf.zeros(n_features, np.float32)

model_coefficients, is_converged, num_iter = tfp.glm.fit_sparse(

    model_matrix=x_,
    response=y_,
    model=model,
    model_coefficients_start=model_coefficients_start,
    l1_regularizer=10,
    l2_regularizer=None,
    maximum_iterations=10,
    maximum_full_sweeps_per_iteration=10,
    tolerance=1e-6,
    learning_rate=None

)

The above code works fine if n_features < n_samples. If not, I get following traceback:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-114-5dea5cef3cb3> in <module>
     25     maximum_full_sweeps_per_iteration=10,
     26     tolerance=1e-6,
---> 27     learning_rate=None
     28 
     29 )

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/glm/proximal_hessian.py in fit_sparse(model_matrix, response, model, model_coefficients_start, tolerance, l1_regularizer, l2_regularizer, maximum_iterations, maximum_full_sweeps_per_iteration, learning_rate, name)
    487         learning_rate=learning_rate,
    488         tolerance=tolerance,
--> 489         name=name)
    490 
    491 

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/optimizer/proximal_hessian_sparse.py in minimize(grad_and_hessian_loss_fn, x_start, tolerance, l1_regularizer, l2_regularizer, maximum_iterations, maximum_full_sweeps_per_iteration, learning_rate, name)
    588             x_start,
    589             tf.zeros([], np.bool, name='converged'),
--> 590             tf.zeros([], np.int32, name='iter'),
    591         ])

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop_v2(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, maximum_iterations, name)
   2476       name=name,
   2477       maximum_iterations=maximum_iterations,
-> 2478       return_same_structure=True)
   2479 
   2480 

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)
   2712                                               list(loop_vars))
   2713       while cond(*loop_vars):
-> 2714         loop_vars = body(*loop_vars)
   2715         if try_to_pack and not isinstance(loop_vars, (list, _basetuple)):
   2716           packed = True

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/optimizer/proximal_hessian_sparse.py in _loop_body(x_start, converged, iter_)
    579           maximum_full_sweeps=maximum_full_sweeps_per_iteration,
    580           tolerance=tolerance,
--> 581           learning_rate=learning_rate)
    582       return x_start, converged, iter_ + 1
    583 

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/optimizer/proximal_hessian_sparse.py in minimize_one_step(gradient_unregularized_loss, hessian_unregularized_loss_outer, hessian_unregularized_loss_middle, x_start, tolerance, l1_regularizer, l2_regularizer, maximum_full_sweeps, learning_rate, name)
    456             tf.zeros(update_shape, dtype=base_dtype, name='x_update'),
    457             tf.zeros(
--> 458                 update_shape, dtype=base_dtype, name='hess_matmul_x_update'),
    459         ])
    460 

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop_v2(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, maximum_iterations, name)
   2476       name=name,
   2477       maximum_iterations=maximum_iterations,
-> 2478       return_same_structure=True)
   2479 
   2480 

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)
   2712                                               list(loop_vars))
   2713       while cond(*loop_vars):
-> 2714         loop_vars = body(*loop_vars)
   2715         if try_to_pack and not isinstance(loop_vars, (list, _basetuple)):
   2716           packed = True

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/optimizer/proximal_hessian_sparse.py in _loop_body(iter_, x_update_diff_norm_sq, x_update, hess_matmul_x_update)
    364       # This is the coordinatewise Newton update if no L1 regularization.
    365       # In above notation, newton_step = -t * (approximation of d/dz|z=0 ULLSC).
--> 366       second_deriv = _hessian_diag_elt_with_l2(coord)
    367       newton_step = -_mul_ignoring_nones(  # pylint: disable=invalid-unary-operand-type
    368           learning_rate, grad_loss_with_l2[..., coord] +

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/optimizer/proximal_hessian_sparse.py in _hessian_diag_elt_with_l2(coord)
    252           axis=-1)
    253       unregularized_component = (
--> 254           hessian_unregularized_loss_middle[..., coord] * inner_square)
    255       l2_component = _mul_or_none(2., l2_regularizer)
    256       return _add_ignoring_nones(unregularized_component, l2_component)

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py in _slice_helper(tensor, slice_spec, var)
    811         ellipsis_mask=ellipsis_mask,
    812         var=var,
--> 813         name=name)
    814 
    815 

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py in strided_slice(input_, begin, end, strides, begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask, var, name)
    977       ellipsis_mask=ellipsis_mask,
    978       new_axis_mask=new_axis_mask,
--> 979       shrink_axis_mask=shrink_axis_mask)
    980 
    981   parent_name = name

~/virtualenvs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_array_ops.py in strided_slice(input, begin, end, strides, begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask, name)
  10370       else:
  10371         message = e.message
> 10372       _six.raise_from(_core._status_to_exception(e.code, message), None)
  10373   # Add nodes to the TensorFlow graph.
  10374   if begin_mask is None:

~/virtualenvs/tf2/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: slice index 1000 of dimension 0 out of bounds. [Op:StridedSlice] name: fit_sparse/minimize/while/minimize_one_step/while/strided_slice/

Although it may not be recommended, is there any reason why the number of features needs to be less than the number of samples?

Start with tensorflow_probability/python/optimizer/proximal_hessian_sparse.py, especially minimize_one_step and the _hessian_diag_elt_with_l2 path shown in the traceback. Reproduce the issue with the supplied n_samples=1000 and n_features=2000 example, then compare it with the n_features < n_samples case. Done means fit_sparse handles the larger feature dimension without the reported StridedSlice error.