# Parallel map of function `f` on the first axis of tensor(s) `elements`. Schematically, `op_vectorized_map()` maps over the first dimension of the provided tensors. If `elements` is a list of tensors, then each of the tensors are required to have the same size first dimension, and they are iterated over together. ```r op_vectorized_map(elements, f) ``` ## Arguments - `elements`: see description - `f`: A function taking either a tensor, or list of tensors. ## Returns A tensor or list of tensors, the result of mapping `f` across `elements.` ## Examples ``` (x <- op_arange(12L) |> op_reshape(c(3, 4))) ``` ``` ## tf.Tensor( ## [[ 1 2 3 4] ## [ 5 6 7 8] ## [ 9 10 11 12]], shape=(3, 4), dtype=int32) ``` ``` x |> op_vectorized_map(\(row) {row + 10}) ``` ``` ## tf.Tensor( ## [[11 12 13 14] ## [15 16 17 18] ## [19 20 21 22]], shape=(3, 4), dtype=int32) ``` ``` list(x, x, x) |> op_vectorized_map(\(rows) Reduce(`+`, rows)) ``` ``` ## tf.Tensor( ## [[ 3 6 9 12] ## [15 18 21 24] ## [27 30 33 36]], shape=(3, 4), dtype=int32) ``` Note that `f` may be traced and compiled. Meaning, the R function may only evaluated once with symbolic tensors if using Jax or TensorFlow backends, and not with eager tensors. See the output from `str()` in these examples: ``` # simplest case, map f over rows of x, # where .x is 1 row of x input <- x output <- op_vectorized_map(input, function(.x) { str(.x) .x + 10 }) ``` ``` ## ``` ``` output ``` ``` ## tf.Tensor( ## [[11 12 13 14] ## [15 16 17 18] ## [19 20 21 22]], shape=(3, 4), dtype=int32) ``` ``` # map f over two tensors simultaneously. Here, # `.x` is a list of two # tensors. The return values from each call of `f(row)` are stacked to form the # final output input <- list(x, x) output <- op_vectorized_map(input, function(.x) { str(.x) .x[[1]] + 10 }) ``` ``` ## List of 2 ## $ : ## $ : ``` ``` output ``` ``` ## tf.Tensor( ## [[11 12 13 14] ## [15 16 17 18] ## [19 20 21 22]], shape=(3, 4), dtype=int32) ``` ``` # same as above, but now returning two tensors in the final output output <- op_vectorized_map(input, function(.x) { str(.x) c(.x1, .x2) %<-% .x list(.x1+10, .x2+20) }) ``` ``` ## List of 2 ## $ : ## $ : ``` ``` output ``` ``` ## [[1]] ## tf.Tensor( ## [[11 12 13 14] ## [15 16 17 18] ## [19 20 21 22]], shape=(3, 4), dtype=int32) ## ## [[2]] ## tf.Tensor( ## [[21 22 23 24] ## [25 26 27 28] ## [29 30 31 32]], shape=(3, 4), dtype=int32) ``` ``` # passing named lists. # WARNING: if passing a named list, the order of elements of `.x` supplied # to `f` is not stable. Only retrieve elements by name. input <- list(name1 = x, name2 = x) output <- op_vectorized_map(input, function(.x) { str(.x) list(outname1 = .x$name1 + 10, outname2 = .x$name2 + 20) }) ``` ``` ## List of 2 ## $ name1: ## $ name2: ``` ``` output ``` ``` ## $outname1 ## tf.Tensor( ## [[11 12 13 14] ## [15 16 17 18] ## [19 20 21 22]], shape=(3, 4), dtype=int32) ## ## $outname2 ## tf.Tensor( ## [[21 22 23 24] ## [25 26 27 28] ## [29 30 31 32]], shape=(3, 4), dtype=int32) ``` ``` # passing a tuple() is equivalent to passing an unnamed list() input <- tuple(x, x) output <- op_vectorized_map(input, function(.x) { str(.x) list(.x[[1]] + 10) }) ``` ``` ## List of 2 ## $ : ## $ : ``` ``` output ``` ``` ## [[1]] ## tf.Tensor( ## [[11 12 13 14] ## [15 16 17 18] ## [19 20 21 22]], shape=(3, 4), dtype=int32) ``` ## Debugging `f` Even in eager contexts, `op_vectorized_map()` may trace `f`. In that case, if you want to eagerly debug `f` (e.g., with `browser()`), you can swap in a manual (slow) implementation of `op_vectorized_map()`. Note this example debug implementation does not handle all the same edge cases as `op_vectorized_map()`, in particular, if `f` returns a structure of multiple tensors. ``` op_vectorized_map_debug <- function(elements, fn) { if (!is.list(elements)) { # `elements` is a single tensor batch_size <- op_shape(elements)[[1]] out <- elements |> op_split(batch_size) |> lapply(fn) |> op_stack() return(out) } # `elements` is a list of tensors batch_size <- elements[[1]] |> op_shape() |> _[[1]] elements |> lapply(\(e) op_split(e, batch_size)) |> zip_lists() |> lapply(fn) |> op_stack() } ``` ## See Also Other core ops: `op_associative_scan()` `op_cast()` `op_cond()` `op_convert_to_numpy()` `op_convert_to_tensor()` `op_custom_gradient()` `op_dtype()` `op_fori_loop()` `op_is_tensor()` `op_map()` `op_rearrange()` `op_scan()` `op_scatter()` `op_scatter_update()` `op_searchsorted()` `op_shape()` `op_slice()` `op_slice_update()` `op_stop_gradient()` `op_subset()` `op_switch()` `op_unstack()` `op_while_loop()` Other ops: `op_abs()` `op_add()` `op_all()` `op_any()` `op_append()` `op_arange()` `op_arccos()` `op_arccosh()` `op_arcsin()` `op_arcsinh()` `op_arctan()` `op_arctan2()` `op_arctanh()` `op_argmax()` `op_argmin()` `op_argpartition()` `op_argsort()` `op_array()` `op_associative_scan()` `op_average()` `op_average_pool()` `op_batch_normalization()` `op_binary_crossentropy()` `op_bincount()` `op_bitwise_and()` `op_bitwise_invert()` `op_bitwise_left_shift()` `op_bitwise_not()` `op_bitwise_or()` `op_bitwise_right_shift()` `op_bitwise_xor()` `op_broadcast_to()` `op_cast()` `op_categorical_crossentropy()` `op_ceil()` `op_celu()` `op_cholesky()` `op_clip()` `op_concatenate()` `op_cond()` `op_conj()` `op_conv()` `op_conv_transpose()` `op_convert_to_numpy()` `op_convert_to_tensor()` `op_copy()` `op_correlate()` `op_cos()` `op_cosh()` `op_count_nonzero()` `op_cross()` `op_ctc_decode()` `op_ctc_loss()` `op_cumprod()` `op_cumsum()` `op_custom_gradient()` `op_depthwise_conv()` `op_det()` `op_diag()` `op_diagflat()` `op_diagonal()` `op_diff()` `op_digitize()` `op_divide()` `op_divide_no_nan()` `op_dot()` `op_dot_product_attention()` `op_dtype()` `op_eig()` `op_eigh()` `op_einsum()` `op_elu()` `op_empty()` `op_equal()` `op_erf()` `op_erfinv()` `op_exp()` `op_exp2()` `op_expand_dims()` `op_expm1()` `op_extract_sequences()` `op_eye()` `op_fft()` `op_fft2()` `op_flip()` `op_floor()` `op_floor_divide()` `op_fori_loop()` `op_full()` `op_full_like()` `op_gelu()` `op_get_item()` `op_glu()` `op_greater()` `op_greater_equal()` `op_hard_shrink()` `op_hard_sigmoid()` `op_hard_silu()` `op_hard_tanh()` `op_histogram()` `op_hstack()` `op_identity()` `op_ifft2()` `op_imag()` `op_image_affine_transform()` `op_image_crop()` `op_image_extract_patches()` `op_image_gaussian_blur()` `op_image_hsv_to_rgb()` `op_image_map_coordinates()` `op_image_pad()` `op_image_perspective_transform()` `op_image_resize()` `op_image_rgb_to_grayscale()` `op_image_rgb_to_hsv()` `op_in_top_k()` `op_inner()` `op_inv()` `op_irfft()` `op_is_tensor()` `op_isclose()` `op_isfinite()` `op_isinf()` `op_isnan()` `op_istft()` `op_leaky_relu()` `op_left_shift()` `op_less()` `op_less_equal()` `op_linspace()` `op_log()` `op_log10()` `op_log1p()` `op_log2()` `op_log_sigmoid()` `op_log_softmax()` `op_logaddexp()` `op_logdet()` `op_logical_and()` `op_logical_not()` `op_logical_or()` `op_logical_xor()` `op_logspace()` `op_logsumexp()` `op_lstsq()` `op_lu_factor()` `op_map()` `op_matmul()` `op_max()` `op_max_pool()` `op_maximum()` `op_mean()` `op_median()` `op_meshgrid()` `op_min()` `op_minimum()` `op_mod()` `op_moments()` `op_moveaxis()` `op_multi_hot()` `op_multiply()` `op_nan_to_num()` `op_ndim()` `op_negative()` `op_nonzero()` `op_norm()` `op_normalize()` `op_not_equal()` `op_one_hot()` `op_ones()` `op_ones_like()` `op_outer()` `op_pad()` `op_polar()` `op_power()` `op_prod()` `op_psnr()` `op_qr()` `op_quantile()` `op_ravel()` `op_real()` `op_rearrange()` `op_reciprocal()` `op_relu()` `op_relu6()` `op_repeat()` `op_reshape()` `op_rfft()` `op_right_shift()` `op_rms_normalization()` `op_roll()` `op_rot90()` `op_round()` `op_rsqrt()` `op_saturate_cast()` `op_scan()` `op_scatter()` `op_scatter_update()` `op_searchsorted()` `op_segment_max()` `op_segment_sum()` `op_select()` `op_selu()` `op_separable_conv()` `op_shape()` `op_sigmoid()` `op_sign()` `op_signbit()` `op_silu()` `op_sin()` `op_sinh()` `op_size()` `op_slice()` `op_slice_update()` `op_slogdet()` `op_soft_shrink()` `op_softmax()` `op_softplus()` `op_softsign()` `op_solve()` `op_solve_triangular()` `op_sort()` `op_sparse_categorical_crossentropy()` `op_sparse_plus()` `op_sparsemax()` `op_split()` `op_sqrt()` `op_square()` `op_squareplus()` `op_squeeze()` `op_stack()` `op_std()` `op_stft()` `op_stop_gradient()` `op_subset()` `op_subtract()` `op_sum()` `op_svd()` `op_swapaxes()` `op_switch()` `op_take()` `op_take_along_axis()` `op_tan()` `op_tanh()` `op_tanh_shrink()` `op_tensordot()` `op_threshold()` `op_tile()` `op_top_k()` `op_trace()` `op_transpose()` `op_tri()` `op_tril()` `op_triu()` `op_trunc()` `op_unravel_index()` `op_unstack()` `op_var()` `op_vdot()` `op_vectorize()` `op_vstack()` `op_where()` `op_while_loop()` `op_zeros()` `op_zeros_like()`

op_vectorized_map function