[CODE CLEAN] torch/testing asserts

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162275
Approved by: https://github.com/mlazos
ghstack dependencies: #164158

.ci/docker/almalinux/Dockerfile

@@ -20,7 +20,7 @@ ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
 
 # cmake-3.18.4 from pip
 RUN yum install -y python3-pip && \
-    python3 -m pip install cmake==3.18.4 && \
+    python3 -mpip install cmake==3.18.4 && \
     ln -s /usr/local/bin/cmake /usr/bin/cmake3
 RUN rm -rf /usr/local/cuda-*
 

.ci/docker/common/install_inductor_benchmark_deps.sh

@@ -25,7 +25,7 @@ function install_torchbench() {
   python install.py --continue_on_fail
 
   echo "Print all dependencies after TorchBench is installed"
-  python -m pip freeze
+  python -mpip freeze
   popd
 
   chown -R jenkins torchbench

.ci/docker/common/install_mkl.sh

@@ -8,8 +8,8 @@ MKLROOT=/opt/intel
 mkdir -p ${MKLROOT}
 pushd /tmp
 
-python3 -m pip install wheel
-python3 -m pip download -d . mkl-static==${MKL_VERSION}
+python3 -mpip install wheel
+python3 -mpip download -d . mkl-static==${MKL_VERSION}
 python3 -m wheel unpack mkl_static-${MKL_VERSION}-py2.py3-none-manylinux1_x86_64.whl
 python3 -m wheel unpack mkl_include-${MKL_VERSION}-py2.py3-none-manylinux1_x86_64.whl
 mv mkl_static-${MKL_VERSION}/mkl_static-${MKL_VERSION}.data/data/lib ${MKLROOT}

.ci/docker/common/install_python.sh

@@ -11,5 +11,5 @@ ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
 python -m venv /var/lib/jenkins/ci_env
 source /var/lib/jenkins/ci_env/bin/activate
 
-python -m pip install --upgrade pip
-python -m pip install -r /opt/requirements-ci.txt
+python -mpip install --upgrade pip
+python -mpip install -r /opt/requirements-ci.txt

.ci/docker/manywheel/Dockerfile_2_28

@@ -14,7 +14,7 @@ ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/op
 
 # cmake-3.18.4 from pip
 RUN yum install -y python3-pip && \
-    python3 -m pip install cmake==3.18.4 && \
+    python3 -mpip install cmake==3.18.4 && \
     ln -s /usr/local/bin/cmake /usr/bin/cmake3
 
 FROM base as openssl
@@ -135,7 +135,7 @@ RUN bash ./patch_libstdc.sh && rm patch_libstdc.sh
 
 # cmake-3.18.4 from pip; force in case cmake3 already exists
 RUN yum install -y python3-pip && \
-    python3 -m pip install cmake==3.18.4 && \
+    python3 -mpip install cmake==3.18.4 && \
     ln -sf /usr/local/bin/cmake /usr/bin/cmake3
 
 FROM cpu_final as cuda_final
@@ -157,7 +157,7 @@ ENV ROCM_PATH /opt/rocm
 # cmake-3.28.4 from pip to get enable_language(HIP)
 # and avoid 3.21.0 cmake+ninja issues with ninja inserting "-Wl,--no-as-needed" in LINK_FLAGS for static linker
 RUN python3 -m pip install --upgrade pip && \
-    python3 -m pip install cmake==3.28.4
+    python3 -mpip install cmake==3.28.4
 # replace the libdrm in /opt/amdgpu with custom amdgpu.ids lookup path
 ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
 RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
@@ -174,7 +174,7 @@ FROM cpu_final as xpu_final
 ENV XPU_DRIVER_TYPE ROLLING
 # cmake-3.28.4 from pip
 RUN python3 -m pip install --upgrade pip && \
-    python3 -m pip install cmake==3.28.4
+    python3 -mpip install cmake==3.28.4
 ADD ./common/install_xpu.sh install_xpu.sh
 ENV XPU_VERSION 2025.2
 RUN bash ./install_xpu.sh && rm install_xpu.sh

.ci/docker/manywheel/Dockerfile_s390x

@@ -113,7 +113,7 @@ RUN dnf install -y \
 RUN env GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=True pip3 install grpcio
 
 # cmake-3.28.0 from pip for onnxruntime
-RUN python3 -m pip install cmake==3.28.0
+RUN python3 -mpip install cmake==3.28.0
 
 ADD ./common/patch_libstdc.sh patch_libstdc.sh
 RUN bash ./patch_libstdc.sh && rm patch_libstdc.sh

.ci/pytorch/build.sh

@@ -288,7 +288,7 @@ else
     # or building non-XLA tests.
     if [[ "$BUILD_ENVIRONMENT" != *rocm*  && "$BUILD_ENVIRONMENT" != *xla* && "$BUILD_ENVIRONMENT" != *riscv64* ]]; then
       # Install numpy-2.0.2 for builds which are backward compatible with 1.X
-      python -m pip install numpy==2.0.2
+      python -mpip install numpy==2.0.2
 
       WERROR=1 python setup.py clean
 

.ci/pytorch/common_utils.sh

@@ -67,13 +67,13 @@ function pip_install_whl() {
     # Loop through each path and install individually
     for path in "${paths[@]}"; do
       echo "Installing $path"
-      python3 -m pip install --no-index --no-deps "$path"
+      python3 -mpip install --no-index --no-deps "$path"
     done
   else
     # Loop through each argument and install individually
     for path in "${args[@]}"; do
       echo "Installing $path"
-      python3 -m pip install --no-index --no-deps "$path"
+      python3 -mpip install --no-index --no-deps "$path"
     done
   fi
 }

.ci/pytorch/macos-test.sh

@@ -182,7 +182,7 @@ checkout_install_torchbench() {
   pip uninstall -y torchao
 
   echo "Print all dependencies after TorchBench is installed"
-  python -m pip freeze
+  python -mpip freeze
 }
 
 torchbench_setup_macos() {
@@ -211,7 +211,7 @@ torchbench_setup_macos() {
 }
 
 pip_benchmark_deps() {
-  python -m pip install --no-input requests cython scikit-learn six
+  python -mpip install --no-input requests cython scikit-learn six
 }
 
 

.ci/pytorch/test.sh

@@ -1434,7 +1434,7 @@ EOF
   # shellcheck source=./common-build.sh
   source "$(dirname "${BASH_SOURCE[0]}")/common-build.sh"
   python -m build --wheel --no-isolation -C--build-option=--bdist-dir="base_bdist_tmp" --outdir "base_dist"
-  python -m pip install base_dist/*.whl
+  python -mpip install base_dist/*.whl
   echo "::endgroup::"
 
   pushd test/forward_backward_compatibility

.ci/wheel/build_wheel.sh

@@ -173,7 +173,7 @@ esac
 PINNED_PACKAGES=(
     "numpy${NUMPY_PINNED_VERSION}"
 )
-python -m venv ~/${desired_python}-build
+python -mvenv ~/${desired_python}-build
 source ~/${desired_python}-build/bin/activate
 retry pip install "${PINNED_PACKAGES[@]}" -r "${pytorch_rootdir}/requirements.txt"
 retry brew install libomp

.github/scripts/prepare_vllm_wheels.sh

@@ -24,7 +24,7 @@ change_wheel_version() {
   local t_version=$4
 
   # Extract the wheel
-  ${PYTHON_EXECUTABLE} -m wheel unpack $wheel
+  ${PYTHON_EXECUTABLE} -mwheel unpack $wheel
 
   mv "${package}-${f_version}" "${package}-${t_version}"
   # Change the version from f_version to t_version in the dist-info dir
@@ -47,7 +47,7 @@ change_wheel_version() {
   popd
 
   # Repack the wheel
-  ${PYTHON_EXECUTABLE} -m wheel pack "${package}-${t_version}"
+  ${PYTHON_EXECUTABLE} -mwheel pack "${package}-${t_version}"
 
   # Clean up
   rm -rf "${package}-${t_version}"
@@ -85,7 +85,7 @@ repackage_wheel() {
 }
 
 # Require to re-package the wheel
-${PYTHON_EXECUTABLE} -m pip install wheel==0.45.1
+${PYTHON_EXECUTABLE} -mpip install wheel==0.45.1
 
 pushd externals/vllm/wheels
 for package in xformers flashinfer-python vllm; do

.github/workflows/_mac-test.yml

@@ -211,7 +211,7 @@ jobs:
             $tool --version
           done
 
-          python3 -m pip install --no-index --no-deps dist/*.whl
+          python3 -mpip install --no-index --no-deps dist/*.whl
 
           set +e
           pushd "${RUNNER_TEMP}"
@@ -222,7 +222,7 @@ jobs:
           popd
 
           if [ "${RC}" -ne 0 ]; then
-            python3 -m pip install --ignore-installed -r "${PIP_REQUIREMENTS_FILE}"
+            python3 -mpip install --ignore-installed -r "${PIP_REQUIREMENTS_FILE}"
           fi
           set -e
 

.github/workflows/_win-test.yml

@@ -204,7 +204,7 @@ jobs:
         run: |
           pushd "${PYTORCH_FINAL_PACKAGE_DIR}"
           # shellcheck disable=SC2046,SC2102
-          python3 -m pip install $(echo *.whl)[opt-einsum,optree] optree==0.13.0
+          python3 -mpip install $(echo *.whl)[opt-einsum,optree] optree==0.13.0
           popd
 
           .ci/pytorch/win-test.sh

.github/workflows/build-vllm-wheel.yml

@@ -126,13 +126,13 @@ jobs:
             "${MANYLINUX_IMAGE}"
           )
 
-          docker exec -t "${container_name}" "${PYTHON_EXECUTABLE}" -m pip install \
+          docker exec -t "${container_name}" "${PYTHON_EXECUTABLE}" -mpip install \
             --pre torch torchvision torchaudio \
             --index-url "https://download.pytorch.org/whl/nightly/${BUILD_DEVICE}"
 
           # I wonder if there is a command to both download and install the wheels
           # in one go
-          docker exec -t "${container_name}" "${PYTHON_EXECUTABLE}" -m pip download \
+          docker exec -t "${container_name}" "${PYTHON_EXECUTABLE}" -mpip download \
             --pre torch torchvision torchaudio \
             --index-url "https://download.pytorch.org/whl/nightly/${BUILD_DEVICE}"
 

.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml

@@ -106,7 +106,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 
@@ -216,7 +216,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 
@@ -326,7 +326,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 
@@ -436,7 +436,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 
@@ -546,7 +546,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 
@@ -656,7 +656,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 
@@ -766,7 +766,7 @@ jobs:
           SMOKE_TEST_PARAMS=""
 
           # shellcheck disable=SC2086
-          python -m venv test_venv
+          python -mvenv test_venv
           source test_venv/bin/activate
           pip install "$PYTORCH_FINAL_PACKAGE_DIR"/*.whl numpy -v
 

Dockerfile

@@ -39,7 +39,7 @@ RUN chmod +x ~/miniconda.sh && \
     bash ~/miniconda.sh -b -p /opt/conda && \
     rm ~/miniconda.sh && \
     /opt/conda/bin/conda install -y python=${PYTHON_VERSION} cmake conda-build pyyaml numpy ipython && \
-    /opt/conda/bin/python -m pip install -r requirements.txt && \
+    /opt/conda/bin/python -mpip install -r requirements.txt && \
     /opt/conda/bin/conda clean -ya
 
 FROM dev-base as submodule-update

tools/linter/adapters/pip_init.py

@@ -76,7 +76,7 @@ def main() -> None:
     if uv and (is_uv_managed_python or not need_user_flag):
         pip_args = [uv, "pip", "install"]
     elif sys.executable:
-        pip_args = [sys.executable, "-m", "pip", "install"]
+        pip_args = [sys.executable, "-mpip", "install"]
     else:
         pip_args = ["pip3", "install"]
 

torch/_inductor/codegen/triton.py

@@ -951,8 +951,7 @@ class TritonCSEVariable(CSEVariable):
         # We'll use this to track which masks the variable needs when used for indirect indexing
         self.mask_vars: OrderedSet[str] = OrderedSet()
         assert dtype is not None, "TritonCSEVariable must have dtype"
-        # TODO: uncomment this and fix the few failures left
-        # assert shape is not None, "TritonCSEVariable must have shape"
+        assert shape is not None, "TritonCSEVariable must have shape"
 
     def update_on_args(self, name, args, kwargs):
         for arg in args:

torch/testing/_internal/common_utils.py

@@ -886,7 +886,8 @@ def cppProfilingFlagsToProfilingMode():
 def enable_profiling_mode_for_profiling_tests():
     old_prof_exec_state = False
     old_prof_mode_state = False
-    assert GRAPH_EXECUTOR
+    if not GRAPH_EXECUTOR:
+        raise AssertionError("GRAPH_EXECUTOR must be set")
     if GRAPH_EXECUTOR == ProfilingMode.PROFILING:
         old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
         old_prof_mode_state = torch._C._get_graph_executor_optimize(True)
@@ -921,7 +922,8 @@ meth_call = torch._C.ScriptMethod.__call__
 def prof_callable(callable, *args, **kwargs):
     if 'profile_and_replay' in kwargs:
         del kwargs['profile_and_replay']
-        assert GRAPH_EXECUTOR
+        if not GRAPH_EXECUTOR:
+            raise AssertionError("GRAPH_EXECUTOR must be set")
         if GRAPH_EXECUTOR == ProfilingMode.PROFILING:
             with enable_profiling_mode_for_profiling_tests():
                 callable(*args, **kwargs)
@@ -1754,7 +1756,8 @@ def serialTest(condition=True):
     """
     # If one apply decorator directly condition will be callable
     # And test will essentially be essentially skipped, which is undesirable
-    assert type(condition) is bool
+    if type(condition) is not bool:
+        raise AssertionError(f"condition must be a bool, got {type(condition)}")
 
     def decorator(fn):
         if has_pytest and condition:
@@ -1811,7 +1814,8 @@ def skipIfLegacyJitExecutor(msg="test doesn't currently work with legacy JIT exe
         if not isinstance(fn, type):
             @wraps(fn)
             def wrapper(*args, **kwargs):
-                assert GRAPH_EXECUTOR
+                if not GRAPH_EXECUTOR:
+                    raise AssertionError("GRAPH_EXECUTOR must be set")
                 if GRAPH_EXECUTOR == ProfilingMode.LEGACY:
                     raise unittest.SkipTest(msg)
                 else:

torch/testing/_internal/distributed/rpc/dist_autograd_test.py

@@ -74,11 +74,13 @@ def _compare_owner_value(context_id, rref, grad):
     grads = dist_autograd.get_gradients(context_id)
     x = grads[rref.local_value()]
     if x.is_sparse:
-        assert grad.is_sparse
+        if not grad.is_sparse:
+            raise AssertionError("grad must be sparse when x is sparse")
         x = x.to_dense()
         grad = grad.to_dense()
     else:
-        assert not grad.is_sparse
+        if grad.is_sparse:
+            raise AssertionError("grad must not be sparse when x is not sparse")
     return torch.equal(x, grad)
 
 
@@ -1940,7 +1942,10 @@ class DistAutogradTest(CommonDistAutogradTest):
         @staticmethod
         @once_differentiable
         def backward(ctx, input):
-            assert DistAutogradTest._test_clean_context_backward_context_id is not None
+            if DistAutogradTest._test_clean_context_backward_context_id is None:
+                raise AssertionError(
+                    "_test_clean_context_backward_context_id must not be None"
+                )
 
             # Release the context to simulate error (use barrier before releasing
             # context to ensure all nodes execute the backward function).
@@ -1950,7 +1955,8 @@ class DistAutogradTest(CommonDistAutogradTest):
             )
 
             # Verify all contexts are cleaned up.
-            assert _all_contexts_cleaned_up()
+            if not _all_contexts_cleaned_up():
+                raise AssertionError("All contexts must be cleaned up")
 
             return input
 
@@ -2050,13 +2056,17 @@ class DistAutogradTest(CommonDistAutogradTest):
         @once_differentiable
         def backward(ctx, input):
             debug_info = dist_autograd._get_debug_info()
-            assert debug_info is not None
+            if debug_info is None:
+                raise AssertionError("debug_info must not be None")
             backward_passes = int(debug_info["num_current_backward_passes"])
 
             # Hard to validate exact numbers because of the distributed nature.
             # We can't use a barrier() here since that would block the single
             # CPU thread available for autograd and can cause deadlocks.
-            assert backward_passes >= 1 and backward_passes <= 4
+            if not (backward_passes >= 1 and backward_passes <= 4):
+                raise AssertionError(
+                    f"backward_passes must be between 1 and 4, got {backward_passes}"
+                )
             return input
 
     @dist_init
@@ -2106,7 +2116,8 @@ class DistAutogradTest(CommonDistAutogradTest):
 
         # Validate information
         debug_info = dist_autograd._get_debug_info()
-        assert debug_info is not None
+        if debug_info is None:
+            raise AssertionError("debug_info must not be None")
         self.assertEqual(0, int(debug_info["num_current_backward_passes"]))
         # only have `num_current_backward_passes` and `num_autograd contexts`
         self.assertTrue(len(debug_info) == 2)

torch/testing/_internal/opinfo/core.py

@@ -1799,10 +1799,22 @@ class ReductionOpInfo(OpInfo):
         assert nan_policy in (None, "propagate", "omit")
 
         # These are mutually exclusive options
-        assert not (result_dtype and promotes_int_to_float)
-        assert not (result_dtype and promotes_int_to_int64)
-        assert not (result_dtype and complex_to_real)
-        assert not (promotes_int_to_float and promotes_int_to_int64)
+        if result_dtype and promotes_int_to_float:
+            raise AssertionError(
+                "result_dtype and promotes_int_to_float are mutually exclusive options"
+            )
+        if result_dtype and promotes_int_to_int64:
+            raise AssertionError(
+                "result_dtype and promotes_int_to_int64 are mutually exclusive options"
+            )
+        if result_dtype and complex_to_real:
+            raise AssertionError(
+                "result_dtype and complex_to_real are mutually exclusive options"
+            )
+        if promotes_int_to_float and promotes_int_to_int64:
+            raise AssertionError(
+                "promotes_int_to_float and promotes_int_to_int64 are mutually exclusive options"
+            )
 
         # Default sample_inputs_func for ReductionOpInfo which augments sample
         # inputs from sample_inputs_reduction with the args and kwargs from
@@ -3069,8 +3081,14 @@ class ForeachFuncInfo(OpInfo):
             # `_getattr_qual` in `OpInfo.__post_init__` which should fail since `_foreach_zero`
             # is not defined at the moment. Thus to skip the qualification, set a similar torch
             # function.
-            assert foreach_method is None
-            assert torch_ref_method is None
+            if foreach_method is not None:
+                raise AssertionError(
+                    "foreach_method must be None when supports_out is False"
+                )
+            if torch_ref_method is not None:
+                raise AssertionError(
+                    "torch_ref_method must be None when supports_out is False"
+                )
             foreach_method = foreach_method_inplace
             torch_ref_method = torch_ref_inplace
 

torch/testing/_internal/opinfo/definitions/sparse.py

@@ -171,7 +171,10 @@ def sample_inputs_sparse_reduction(
                 dtype=inp.dtype,
                 device=inp.device,
             )
-            assert not inp.is_coalesced()
+            if inp.is_coalesced():
+                raise AssertionError(
+                    "Expected uncoalesced sparse tensor, but tensor is coalesced"
+                )
             yield SampleInput(
                 inp.requires_grad_(requires_grad),
                 args=sample_input.args,

torch/testing/_internal/opinfo/utils.py

@@ -230,7 +230,10 @@ def reference_reduction_numpy(f, supports_keepdims=True):
         if "mask" in keys:
             mask = kwargs.pop("mask")
             if mask is not None:
-                assert mask.layout == torch.strided
+                if mask.layout != torch.strided:
+                    raise AssertionError(
+                        f"mask must have strided layout, got {mask.layout}"
+                    )
                 kwargs["where"] = mask.cpu().numpy()
 
         if "identity" in keys:

torch/testing/_internal/subclasses.py

@@ -38,11 +38,16 @@ class WrapperSubclass(torch.Tensor):
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-        assert meta is None
+        if meta is not None:
+            raise AssertionError(f"Expected meta to be None, got {meta}")
         a = inner_tensors["a"]
         if is_fake(a):
-            assert outer_size is not None
-            assert outer_stride is not None
+            if outer_size is None:
+                raise AssertionError("outer_size must not be None when tensor is fake")
+            if outer_stride is None:
+                raise AssertionError(
+                    "outer_stride must not be None when tensor is fake"
+                )
         return WrapperSubclass(a, outer_size, outer_stride)
 
     @classmethod

torch/testing/_internal/triton_utils.py

@@ -1032,7 +1032,10 @@ if has_triton():
                     tensor.element_size(),
                 )
             else:
-                assert len(block_sizes) == 2
+                if len(block_sizes) != 2:
+                    raise AssertionError(
+                        f"block_sizes must have exactly 2 elements for 2D TMA descriptor, got {len(block_sizes)}"
+                    )
                 return triton.tools.experimental_descriptor.create_2d_tma_descriptor(
                     tensor.data_ptr(),
                     tensor.size(0),

torch/testing/_internal/two_tensor.py

@@ -15,12 +15,17 @@ class TwoTensor(torch.Tensor):
         if outer_stride is None:
             outer_stride = a.stride()
 
-        assert (
+        if not (
             a.device == b.device
             and a.layout == b.layout
             and a.requires_grad == b.requires_grad
             and a.dtype == b.dtype
-        )
+        ):
+            raise AssertionError(
+                f"Tensors a and b must have matching device, layout, requires_grad, and dtype. "
+                f"Got a: device={a.device}, layout={a.layout}, requires_grad={a.requires_grad}, dtype={a.dtype}; "
+                f"b: device={b.device}, layout={b.layout}, requires_grad={b.requires_grad}, dtype={b.dtype}"
+            )
         # I guess it would be more accurate to represent the shape as torch.cat(a, b).shape
         shape = outer_size
         kwargs = {}
@@ -32,9 +37,18 @@ class TwoTensor(torch.Tensor):
         kwargs["dtype"] = a.dtype
         out = torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
 
-        assert a.shape == b.shape
-        assert a.stride() == b.stride()
-        assert a.storage_offset() == b.storage_offset()
+        if a.shape != b.shape:
+            raise AssertionError(
+                f"Tensors a and b must have the same shape. Got a.shape={a.shape}, b.shape={b.shape}"
+            )
+        if a.stride() != b.stride():
+            raise AssertionError(
+                f"Tensors a and b must have the same stride. Got a.stride()={a.stride()}, b.stride()={b.stride()}"
+            )
+        if a.storage_offset() != b.storage_offset():
+            raise AssertionError(
+                f"Tensors a and b must have the same storage_offset. Got a.storage_offset()={a.storage_offset()}, b.storage_offset()={b.storage_offset()}"
+            )
         return out
 
     @torch._disable_dynamo
@@ -53,11 +67,18 @@ class TwoTensor(torch.Tensor):
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
-        assert meta is None
+        if meta is not None:
+            raise AssertionError(f"Expected meta to be None, got {meta}")
         a, b = inner_tensors["a"], inner_tensors["b"]
         if type(a) is torch.Tensor:
-            assert outer_size is not None
-            assert outer_stride is not None
+            if outer_size is None:
+                raise AssertionError(
+                    "outer_size must not be None when inner tensor is a torch.Tensor"
+                )
+            if outer_stride is None:
+                raise AssertionError(
+                    "outer_stride must not be None when inner tensor is a torch.Tensor"
+                )
         return TwoTensor(a, b, outer_size, outer_stride)
 
     @classmethod

torch/utils/collect_env.py

@@ -640,9 +640,9 @@ def get_pip_packages(run_lambda, patterns=None):
 
     os.environ["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
     # People generally have pip as `pip` or `pip3`
-    # But here it is invoked as `python -m pip`
+    # But here it is invoked as `python -mpip`
     out = run_and_read_all(
-        run_lambda, [sys.executable, "-m", "pip", "list", "--format=freeze"]
+        run_lambda, [sys.executable, "-mpip", "list", "--format=freeze"]
     )
     if out is None:
         return pip_version, out