Lines Matching defs:device
97 def random_nt_noncontiguous_pair(ragged_sizes, device="cpu", dtype=torch.float16): argument
143 device, argument
187 dims, device=None, dtype=None, layout=torch.strided, requires_grad=False argument
852 def random_nt_pair(self, device, dtype, num_tensors, max_dims): argument
872 def test_detach(self, device, dtype): argument
900 def test_unbind_noncontiguous(self, device, dtype): argument
913 def test_to_then_from_padded_tensor_no_transform0213(self, device, dtype): argument
930 def test_layer_norm(self, device, dtype): argument
991 def test_layer_norm_breaking(self, device, dtype): argument
1018 def test_embedding(self, device, layout): argument
1040 def test_masked_fill(self, device, dtype): argument
1051 def test_to_padded_tensor_simple(self, device, dtype): argument
1070 def test_to_padded_tensor_output_size(self, device, dtype): argument
1094 def test_to_padded_tensor_dim2(self, device, dtype): argument
1112 def test_to_padded_tensor_dim3(self, device, dtype): argument
1130 def test_to_padded_tensor_dim4(self, device, dtype): argument
1153 def test_to_padded_tensor_noncontiguous(self, device, dtype): argument
1170 def test_device_checks(self, device): argument
1176 def test_nested_tensor_indexing(self, device, dtype): argument
1245 def test_activations(self, device, func): argument
1260 def test_binary_ops_with_scalar(self, device, func): argument
1274 def test_nested_tensor_chunk(self, device, dtype): argument
1337 def test_nested_tensor_split_with_sizes(self, device, dtype): argument
1397 def test_nested_tensor_indexing_noncontiguous(self, device, dtype): argument
1410 def test_nested_tensor_add(self, device, dtype, transpose): argument
1430 def test_nested_tensor_sub(self, device, dtype, transpose): argument
1450 def test_nested_tensor_dense_elementwise(self, device, dtype, embedding_dim): argument
1479 def test_nested_tensor_mul(self, device, dtype): argument
1515 def test_nested_tensor_div(self, device, dtype): argument
1557 def test_nested_tensor_add_in_place(self, device, dtype): argument
1568 def test_nested_tensor_mul_in_place(self, device, dtype): argument
1607 def test_nested_tensor_sum_dim(self, device, dtype): argument
1610 def test_sum(device, dtype, ntensors, max_sizes, dim, keepdim=True): argument
1659 def test_contiguous(self, device, dtype): argument
1682 def test_clone(self, device, dtype): argument
1703 def test_dropout(self, device, dtype, layout): argument
1761 def test_dropout_noncontiguous(self, device, dtype): argument
1775 def test_softmax(self, device, dtype): argument
1814 def test_softmax_noncontiguous(self, device, dtype): argument
1823 def _test_bmm(self, device, dtype): argument
1952 def test_bmm_cuda(self, device, dtype): argument
1958 def test_bmm_cpu(self, device, dtype): argument
1963 def test_bmm_noncontiguous(self, device, dtype): argument
1976 def test_matmul_with_bmm_path(self, device, dtype): argument
2018 def test_matmul(self, device, dtype): argument
2180 def test_matmul_nt_with_broadcasted_t(self, device, dtype): argument
2193 def test_matmul_noncontiguous(self, device, dtype): argument
2206 def test_linear(self, device, dtype): argument
2264 def test_linear_noncontiguous(self, device, dtype): argument
2276 def test_to_padded_tensor_zero_numel_errors(self, device, dtype): argument
2288 def test_transpose(self, device, dtype): argument
2312 def test_squeeze_unsqueeze(self, device, dtype): argument
2372 def test_transpose_inference_mode_interaction(self, device, dtype): argument
2392 def test_view(self, device, dtype): argument
2447 def test_view_inference_mode_interaction(self, device, dtype): argument
2470 def test_reshape(self, device, dtype): argument
2524 def test_nested_masked_select(self, device): argument
2572 def test_narrow(self, device, dtype): argument
2600 def test_scaled_dot_product_attention(self, device, input_dim): argument
2677 def test_empty_like(self, device, dtype): argument
2774 def test_as_nested_tensor_propagates_gradients(self, device): argument
2792 def test_nested_tensor_generates_leaf(self, device): argument
2812 def test_set_requires_grad_from_list(self, device): argument
2817 def test_set_requires_grad_from_mask(self, device): argument
2822 def test_backward_for_add_op(self, device): argument
2838 def test_backward_for_sub_op(self, device): argument
2855 def test_backward_sub_strided(self, device): argument
2872 def test_backward_add_strided(self, device): argument
2890 def test_nested_tensor_to_padded_tensor(self, device): argument
2907 def test_nested_tensor_from_mask_and_to_padded(self, device): argument
2928 def test_nested_tensor_from_padded(self, device): argument
2944 def test_nested_tensor_from_padded_fused(self, device): argument
2960 def test_nested_tensor_from_list(self, device): argument
2992 def test_nested_tensor_bmm_gradcheck(self, device): argument
3007 def test_nested_tensor_bmm_backward(self, device): argument
3030 def test_nested_tensor_matmul_gradcheck(self, device): argument
3045 def test_nested_tensor_matmul_backward(self, device): argument
3068 def test_nested_tensor_transpose_gradcheck(self, device): argument
3080 def test_nested_tensor_transpose_backward(self, device): argument
3096 def test_nested_tensor_reshape_gradcheck(self, device): argument
3122 def test_nested_tensor_squeeze_backward(self, device): argument
3138 def test_nested_tensor_squeeze_gradcheck(self, device): argument
3153 def test_nested_tensor_unsqueeze_backward(self, device): argument
3169 def test_nested_tensor_unsqueeze_gradcheck(self, device): argument
3180 def test_nested_tensor_linear(self, device): argument
3203 def test_nested_tensor_linear_plus_transpose(self, device): argument
3227 def test_nested_tensor_softmax(self, device): argument
3242 def test_nested_tensor_linear_backward(self, device): argument
3262 def test_values_grad_with_broadcast(self, device): argument
3275 def test_to_buffer_series_ops_grad_with_broadcast(self, device): argument
3289 def test_unbind_flow_through(self, device): argument
3305 def test_split_with_sizes_flow_through(self, device): argument
3321 def test_indexing_backward(self, device): argument
3334 def test_masked_fill_backward(self, device): argument
3349 def test_gelu_backward(self, device): argument
3362 def test_relu_backward(self, device): argument
3375 def test_selu_backward(self, device): argument
3388 def test_abs_backward(self, device): argument
3403 def test_layer_norm_backward_edge_case(self, device): argument
3415 def test_accumulate_grad_different_strides(self, device): argument
3431 def test_layer_norm_backward(self, device, size): argument
3457 def test_layer_norm_backward_5d(self, device, size): argument
3523 def _get_list_for_jagged_tensor(self, nested_size, device, requires_grad=True): argument
3622 def test_tensor_attributes(self, device): argument
3654 def test_linear(self, device, nt_dim): argument
3682 def test_unary_pointwise(self, device): argument
3694 def test_unary_pointwise_transposed_inputs(self, device): argument
3730 def test_binary_pointwise(self, device): argument
3756 def test_binary_pointwise_transposed(self, device): argument
3796 def test_split(self, device): argument
3823 def test_split_with_sizes(self, device): argument
3849 def test_softmax(self, device): argument
3883 def test_views_inherit_ragged_dim(self, device): argument
3903 def test_view_ragged_idx_not_one(self, device): argument
3912 def test_unsafe_view(self, device): argument
3930 def test_reshape_decomp(self, device, requires_grad): argument
3968 def test_flatten_decomp(self, device): argument
3981 def test_chunk(self, device): argument
4055 def test_squeeze(self, device): argument
4092 def test_binary_pointwise_broadcasting(self, device): argument
4121 def test_threshold_backward(self, device): argument
4149 self, device, dtype, keepdim, requires_grad, components_require_grad, func argument
4236 device, argument
4308 self, device, dtype, keepdim, requires_grad, components_require_grad, func argument
4350 self, device, dtype, requires_grad, components_require_grad argument
4390 self, device, dtype, requires_grad, components_require_grad argument
4421 self, device, dtype, requires_grad, components_require_grad argument
4475 device, argument
4513 device, argument
4558 device, argument
4619 device, argument
4666 self, device, dtype, keepdim, requires_grad, components_require_grad, func argument
4736 device, argument
4788 device, argument
4827 device, argument
4870 self, device, dtype, keepdim, requires_grad, components_require_grad, func argument
4909 device, argument
4960 device, argument
5010 device, argument
5064 device, argument
5100 device, argument
5136 def test_serialization(self, device, dtype, requires_grad, weights_only): argument
5162 def test_pin_memory(self, device): argument
5178 device, argument
5237 self, device, dtype, requires_grad, components_require_grad argument
5284 self, device, dtype, components_require_grad argument
5333 def test_jagged_layout_construction_with_pinned_memory(self, device): argument
5367 self, device, dtype, requires_grad, values_is_view argument
5416 def test_nested_tensor_from_jagged(self, device, dtype, pass_min_max): argument
5525 def test_nested_tensor_from_jagged_fx_trace(self, device): argument
5548 self, device, dtype, dim, layout, requires_grad, contiguous argument
5639 def test_device_dtype_transfer_updates_offsets(self, device, dtype): argument
5654 def test_unbind(self, device): argument
5665 def test_unbind_transpose(self, device, ragged_idx): argument
5679 def test_unbind_transpose_ragged_idx_last_dim(self, device): argument
5691 def test_unbind_lengths(self, device): argument
5708 def test_unbind_lengths_ragged_idx_1(self, device): argument
5727 def test_unbind_lengths_ragged_idx_equals_2_bad_dim(self, device): argument
5742 def test_unbind_lengths_ragged_idx_2(self, device): argument
5761 def test_unbind_lengths_ragged_idx_3(self, device): argument
5783 def test_unbind_lengths_ragged_idx_0(self, device): argument
5802 def test_narrow(self, device): argument
5828 def test_njt_cat(self, device): argument
5849 def test_is_contiguous(self, device): argument
5931 def test_noncontiguous_pointwise(self, device): argument
5958 def test_permute(self, device): argument
5989 def test_to_dtype(self, device): argument
6005 def test_to_copy(self, device): argument
6023 def test_copy_(self, device): argument
6096 def test_is_same_size(self, device): argument
6124 def test_specialize_dynamic_shape(self, device): argument
6146 def test_specialize_dynamic_shape_recompile(self, device): argument
6194 def test_sdpa(self, device, dtype): argument
6436 def test_sdpa_compile(self, device, dtype): argument
6524 def test_sdpa_with_constant_sequence_length(self, device, dtype): argument
6564 def test_sdpa_with_packed_in_proj(self, device, dtype): argument
6622 def test_sdpa_backwards(self, device, dtype): argument
6648 def test_sdpa_autocast(self, device): argument
6727 def test_sdpa_flop_counter(self, device): argument
6752 def test_nested_tensor_activation_checkpoint(self, device): argument
6790 def test_dummy_mha_with_nt(self, device, use_legacy_api): argument
6881 def test_apply_(self, device, dtype): argument
6910 def test_jagged_padded_dense_conversion_kernels(self, device, dtype): argument
7005 def test_compile_preserves_metadata_cache(self, device, dtype): argument
7036 def test_compile_with_dynamic_max_seq_len(self, device, dtype): argument
7073 def test_compile_with_dynamic_min_seq_len(self, device, dtype): argument
7110 def test_compile_with_propagated_dynamic_max_seq_len(self, device, dtype): argument
7145 def test_unbind_backward(self, device, dtype): argument
7304 def test_forward(self, device, dtype, op): argument
7316 def test_backward(self, device, dtype, op): argument
7345 def test_compile_forward(self, device, dtype, op): argument
7372 def test_compile_backward(self, device, dtype, op): argument