FM / python

/usr/local/lib64/python3.6/site-packages/caffe2/python/

/usr/local/lib64/python3.6/site-packages/caffe2/python

Name	Size	Mode	Actions
docs/	-	0755	rm
examples/	-	0755	rm
fakelowp/	-	0755	rm
helpers/	-	0755	rm
ideep/	-	0755	rm
layers/	-	0755	rm
mint/	-	0755	rm
mkl/	-	0755	rm
modeling/	-	0755	rm
models/	-	0755	rm
onnx/	-	0755	rm
operator_test/	-	0755	rm
predictor/	-	0755	rm
rnn/	-	0755	rm
serialized_test/	-	0755	rm
test/	-	0755	rm
trt/	-	0755	rm
__pycache__/	-	0755	rm
allcompare_test.py	2255	0644	edit dl rm
attention.py	12359	0644	edit dl rm
benchmark_generator.py	4912	0644	edit dl rm
binarysize.py	5521	0644	edit dl rm
brew.py	4762	0644	edit dl rm
brew_test.py	11739	0644	edit dl rm
build.py	153	0644	edit dl rm
cached_reader.py	4394	0644	edit dl rm
caffe2_pybind11_state.cpython-36m-x86_64-linux-gnu.so	48299712	0755	edit dl rm
caffe2_pybind11_state_gpu.cpython-36m-x86_64-linux-gnu.so	49048144	0755	edit dl rm
caffe_translator.py	35227	0644	edit dl rm
caffe_translator_test.py	3553	0644	edit dl rm
checkpoint.py	32101	0644	edit dl rm
checkpoint_test.py	13405	0644	edit dl rm
cnn.py	7626	0644	edit dl rm
context.py	2841	0644	edit dl rm
context_test.py	1792	0644	edit dl rm
control.py	19309	0644	edit dl rm
control_ops_grad.py	28893	0644	edit dl rm
control_ops_grad_test.py	1752	0644	edit dl rm
control_ops_util.py	10863	0644	edit dl rm
control_test.py	12276	0644	edit dl rm
convert.py	55	0644	edit dl rm
convert_test.py	201	0644	edit dl rm
convnet_benchmarks.py	20533	0644	edit dl rm
convnet_benchmarks_test.py	839	0644	edit dl rm
core.py	119400	0644	edit dl rm
core_gradients_test.py	38022	0644	edit dl rm
core_test.py	47683	0644	edit dl rm
crf.py	13250	0644	edit dl rm
crf_predict.py	1159	0644	edit dl rm
crf_viterbi_test.py	1663	0644	edit dl rm
dataio.py	23532	0644	edit dl rm
dataio_test.py	17575	0644	edit dl rm
dataset.py	12886	0644	edit dl rm
data_parallel_model.py	83100	0644	edit dl rm
data_parallel_model_test.py	56145	0644	edit dl rm
data_workers.py	15941	0644	edit dl rm
data_workers_test.py	6561	0644	edit dl rm
db_file_reader.py	6608	0644	edit dl rm
db_test.py	1110	0644	edit dl rm
device_checker.py	5157	0644	edit dl rm
dyndep.py	1533	0644	edit dl rm
embedding_generation_benchmark.py	5256	0644	edit dl rm
experiment_util.py	3625	0644	edit dl rm
extension_loader.py	744	0644	edit dl rm
fakefp16_transform_lib.py	322	0644	edit dl rm
filler_test.py	748	0644	edit dl rm
functional.py	4415	0644	edit dl rm
functional_test.py	4204	0644	edit dl rm
fused_8bit_rowwise_conversion_ops_test.py	3945	0644	edit dl rm
gradient_checker.py	15377	0644	edit dl rm
gradient_check_test.py	20729	0644	edit dl rm
gru_cell.py	5129	0644	edit dl rm
hip_test_util.py	405	0644	edit dl rm
hsm_util.py	2259	0644	edit dl rm
hypothesis_test.py	105762	0644	edit dl rm
hypothesis_test_util.py	26853	0644	edit dl rm
ideep_test_util.py	998	0644	edit dl rm
layers_test.py	92931	0644	edit dl rm
layer_model_helper.py	29340	0644	edit dl rm
layer_model_instantiator.py	3935	0644	edit dl rm
layer_parameter_sharing_test.py	9148	0644	edit dl rm
layer_test_util.py	4875	0644	edit dl rm
lazy.py	277	0644	edit dl rm
lazy_dyndep.py	2562	0644	edit dl rm
lazy_dyndep_test.py	3914	0644	edit dl rm
lengths_reducer_fused_8bit_rowwise_ops_test.py	7575	0644	edit dl rm
lengths_reducer_rowwise_8bit_ops_test.py	5710	0644	edit dl rm
lstm_benchmark.py	10649	0644	edit dl rm
memonger.py	34041	0644	edit dl rm
memonger_test.py	36910	0644	edit dl rm
mkl_test_util.py	1142	0644	edit dl rm
model_device_test.py	4777	0644	edit dl rm
model_helper.py	23492	0644	edit dl rm
model_helper_test.py	2336	0644	edit dl rm
modifier_context.py	1772	0644	edit dl rm
muji.py	8131	0644	edit dl rm
muji_test.py	3058	0644	edit dl rm
net_builder.py	27679	0644	edit dl rm
net_builder_test.py	11382	0644	edit dl rm
net_drawer.py	14264	0644	edit dl rm
net_printer.py	12704	0644	edit dl rm
net_printer_test.py	3190	0644	edit dl rm
nomnigraph.py	4216	0644	edit dl rm
nomnigraph_test.py	15427	0644	edit dl rm
nomnigraph_transformations.py	3787	0644	edit dl rm
nomnigraph_transformations_test.py	5767	0644	edit dl rm
normalizer.py	1411	0644	edit dl rm
normalizer_context.py	1007	0644	edit dl rm
normalizer_test.py	487	0644	edit dl rm
numa_benchmark.py	2230	0644	edit dl rm
numa_test.py	1663	0644	edit dl rm
observer_test.py	5316	0644	edit dl rm
operator_fp_exceptions_test.py	1248	0644	edit dl rm
optimizer.py	78813	0644	edit dl rm
optimizer_context.py	1462	0644	edit dl rm
optimizer_test.py	30705	0644	edit dl rm
optimizer_test_util.py	9187	0644	edit dl rm
parallelize_bmuf_distributed_test.py	9908	0644	edit dl rm
parallel_workers.py	7682	0644	edit dl rm
parallel_workers_test.py	3501	0644	edit dl rm
pipeline.py	17283	0644	edit dl rm
pipeline_test.py	2542	0644	edit dl rm
predictor_constants.py	198	0644	edit dl rm
python_op_test.py	9169	0644	edit dl rm
queue_util.py	4459	0644	edit dl rm
record_queue.py	4453	0644	edit dl rm
recurrent.py	13297	0644	edit dl rm
regularizer.py	21120	0644	edit dl rm
regularizer_context.py	1013	0644	edit dl rm
regularizer_test.py	10266	0644	edit dl rm
rnn_cell.py	68233	0644	edit dl rm
schema.py	45621	0644	edit dl rm
schema_test.py	15754	0644	edit dl rm
scope.py	3623	0644	edit dl rm
scope_test.py	5249	0644	edit dl rm
session.py	7642	0644	edit dl rm
session_test.py	2078	0644	edit dl rm
sparse_to_dense_mask_test.py	6565	0644	edit dl rm
sparse_to_dense_test.py	3556	0644	edit dl rm
task.py	24274	0644	edit dl rm
task_test.py	870	0644	edit dl rm
test_util.py	3524	0644	edit dl rm
text_file_reader.py	1990	0644	edit dl rm
timeout_guard.py	4054	0644	edit dl rm
toy_regression_test.py	2822	0644	edit dl rm
transformations.py	1832	0644	edit dl rm
transformations_test.py	11960	0644	edit dl rm
tt_core.py	9349	0644	edit dl rm
tt_core_test.py	2518	0644	edit dl rm
utils.py	14181	0644	edit dl rm
utils_test.py	1399	0644	edit dl rm
visualize.py	6315	0644	edit dl rm
workspace.py	25263	0644	edit dl rm
workspace_test.py	34844	0644	edit dl rm
_import_c_extension.py	2250	0644	edit dl rm
__init__.py	3925	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/workspace_test.py (34844B)

import os
import shutil
import tempfile
import unittest
from collections import namedtuple

import caffe2.python.hypothesis_test_util as htu
import hypothesis.strategies as st
import numpy as np
import torch
from caffe2.proto import caffe2_pb2
from caffe2.python import core, test_util, workspace, model_helper, brew
from hypothesis import given, settings

class TestWorkspace(unittest.TestCase):
    def setUp(self):
        self.net = core.Net("test-net")
        self.testblob_ref = self.net.ConstantFill(
            [], "testblob", shape=[1, 2, 3, 4], value=1.0
        )
        workspace.ResetWorkspace()

def testRootFolder(self):
        self.assertEqual(workspace.ResetWorkspace(), True)
        self.assertEqual(workspace.RootFolder(), ".")
        self.assertEqual(workspace.ResetWorkspace("/tmp/caffe-workspace-test"), True)
        self.assertEqual(workspace.RootFolder(), "/tmp/caffe-workspace-test")

def testWorkspaceHasBlobWithNonexistingName(self):
        self.assertEqual(workspace.HasBlob("non-existing"), False)

def testRunOperatorOnce(self):
        self.assertEqual(
            workspace.RunOperatorOnce(self.net.Proto().op[0].SerializeToString()), True
        )
        self.assertEqual(workspace.HasBlob("testblob"), True)
        blobs = workspace.Blobs()
        self.assertEqual(len(blobs), 1)
        self.assertEqual(blobs[0], "testblob")

def testGetOperatorCost(self):
        op = core.CreateOperator(
            "Conv2D",
            ["X", "W"],
            ["Y"],
            stride_h=1,
            stride_w=1,
            pad_t=1,
            pad_l=1,
            pad_b=1,
            pad_r=1,
            kernel=3,
        )
        X = np.zeros((1, 8, 8, 8))
        W = np.zeros((1, 1, 3, 3))
        workspace.FeedBlob("X", X)
        workspace.FeedBlob("W", W)
        op_cost = workspace.GetOperatorCost(op.SerializeToString(), ["X", "W"])
        self.assertTupleEqual(
            op_cost,
            namedtuple("Cost", ["flops", "bytes_written", "bytes_read"])(
                1152, 256, 4168
            ),
        )

def testRunNetOnce(self):
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        self.assertEqual(workspace.HasBlob("testblob"), True)

def testCurrentWorkspaceWrapper(self):
        self.assertNotIn("testblob", workspace.C.Workspace.current.blobs)
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        self.assertEqual(workspace.HasBlob("testblob"), True)
        self.assertIn("testblob", workspace.C.Workspace.current.blobs)
        workspace.ResetWorkspace()
        self.assertNotIn("testblob", workspace.C.Workspace.current.blobs)

def testRunPlan(self):
        plan = core.Plan("test-plan")
        plan.AddStep(core.ExecutionStep("test-step", self.net))
        self.assertEqual(workspace.RunPlan(plan.Proto().SerializeToString()), True)
        self.assertEqual(workspace.HasBlob("testblob"), True)

def testRunPlanInBackground(self):
        plan = core.Plan("test-plan")
        plan.AddStep(core.ExecutionStep("test-step", self.net))
        background_plan = workspace.RunPlanInBackground(plan)
        while not background_plan.is_done():
            pass
        self.assertEqual(background_plan.is_succeeded(), True)
        self.assertEqual(workspace.HasBlob("testblob"), True)

def testConstructPlanFromSteps(self):
        step = core.ExecutionStep("test-step-as-plan", self.net)
        self.assertEqual(workspace.RunPlan(step), True)
        self.assertEqual(workspace.HasBlob("testblob"), True)

def testResetWorkspace(self):
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        self.assertEqual(workspace.HasBlob("testblob"), True)
        self.assertEqual(workspace.ResetWorkspace(), True)
        self.assertEqual(workspace.HasBlob("testblob"), False)

def testTensorAccess(self):
        ws = workspace.C.Workspace()

""" test in-place modification """
        ws.create_blob("tensor").feed(np.array([1.1, 1.2, 1.3]))
        tensor = ws.blobs["tensor"].tensor()
        tensor.data[0] = 3.3
        val = np.array([3.3, 1.2, 1.3])
        np.testing.assert_array_equal(tensor.data, val)
        np.testing.assert_array_equal(ws.blobs["tensor"].fetch(), val)

""" test in-place initialization """
        tensor.init([2, 3], core.DataType.INT32)
        for x in range(2):
            for y in range(3):
                tensor.data[x, y] = 0
        tensor.data[1, 1] = 100
        val = np.zeros([2, 3], dtype=np.int32)
        val[1, 1] = 100
        np.testing.assert_array_equal(tensor.data, val)
        np.testing.assert_array_equal(ws.blobs["tensor"].fetch(), val)

""" strings cannot be initialized from python """
        with self.assertRaises(RuntimeError):
            tensor.init([3, 4], core.DataType.STRING)

""" feed (copy) data into tensor """
        val = np.array([[b"abc", b"def"], [b"ghi", b"jkl"]], dtype=np.object)
        tensor.feed(val)
        self.assertEquals(tensor.data[0, 0], b"abc")
        np.testing.assert_array_equal(ws.blobs["tensor"].fetch(), val)

val = np.array([1.1, 10.2])
        tensor.feed(val)
        val[0] = 5.2
        self.assertEquals(tensor.data[0], 1.1)

""" fetch (copy) data from tensor """
        val = np.array([1.1, 1.2])
        tensor.feed(val)
        val2 = tensor.fetch()
        tensor.data[0] = 5.2
        val3 = tensor.fetch()
        np.testing.assert_array_equal(val, val2)
        self.assertEquals(val3[0], 5.2)

def testFetchFeedBlob(self):
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        fetched = workspace.FetchBlob("testblob")
        # check if fetched is correct.
        self.assertEqual(fetched.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched, 1.0)
        fetched[:] = 2.0
        self.assertEqual(workspace.FeedBlob("testblob", fetched), True)
        fetched_again = workspace.FetchBlob("testblob")
        self.assertEqual(fetched_again.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched_again, 2.0)

def testFetchFeedBlobViaBlobReference(self):
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        fetched = workspace.FetchBlob(self.testblob_ref)
        # check if fetched is correct.
        self.assertEqual(fetched.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched, 1.0)
        fetched[:] = 2.0
        self.assertEqual(workspace.FeedBlob(self.testblob_ref, fetched), True)
        fetched_again = workspace.FetchBlob("testblob")  # fetch by name now
        self.assertEqual(fetched_again.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched_again, 2.0)

def testFetchFeedBlobTypes(self):
        for dtype in [
            np.float16,
            np.float32,
            np.float64,
            np.bool,
            np.int8,
            np.int16,
            np.int32,
            np.int64,
            np.uint8,
            np.uint16,
        ]:
            try:
                rng = np.iinfo(dtype).max * 2
            except ValueError:
                rng = 1000
            data = ((np.random.rand(2, 3, 4) - 0.5) * rng).astype(dtype)
            self.assertEqual(workspace.FeedBlob("testblob_types", data), True)
            fetched_back = workspace.FetchBlob("testblob_types")
            self.assertEqual(fetched_back.shape, (2, 3, 4))
            self.assertEqual(fetched_back.dtype, dtype)
            np.testing.assert_array_equal(fetched_back, data)

def testFetchFeedBlobBool(self):
        """Special case for bool to ensure coverage of both true and false."""
        data = np.zeros((2, 3, 4)).astype(np.bool)
        data.flat[::2] = True
        self.assertEqual(workspace.FeedBlob("testblob_types", data), True)
        fetched_back = workspace.FetchBlob("testblob_types")
        self.assertEqual(fetched_back.shape, (2, 3, 4))
        self.assertEqual(fetched_back.dtype, np.bool)
        np.testing.assert_array_equal(fetched_back, data)

def testGetBlobSizeBytes(self):
        for dtype in [
            np.float16,
            np.float32,
            np.float64,
            np.bool,
            np.int8,
            np.int16,
            np.int32,
            np.int64,
            np.uint8,
            np.uint16,
        ]:
            data = np.random.randn(2, 3).astype(dtype)
            self.assertTrue(workspace.FeedBlob("testblob_sizeBytes", data), True)
            self.assertEqual(
                workspace.GetBlobSizeBytes("testblob_sizeBytes"),
                6 * np.dtype(dtype).itemsize,
            )
        strs1 = np.array([b"Hello World!", b"abcd"])
        strs2 = np.array([b"element1", b"element2"])
        strs1_len, strs2_len = 0, 0
        for str in strs1:
            strs1_len += len(str)
        for str in strs2:
            strs2_len += len(str)
        self.assertTrue(workspace.FeedBlob("testblob_str1", strs1), True)
        self.assertTrue(workspace.FeedBlob("testblob_str2", strs2), True)
        # size of blob "testblob_str1" = size_str1 * meta_.itemsize() + strs1_len
        # size of blob "testblob_str2" = size_str2 * meta_.itemsize() + strs2_len
        self.assertEqual(
            workspace.GetBlobSizeBytes("testblob_str1")
            - workspace.GetBlobSizeBytes("testblob_str2"),
            strs1_len - strs2_len,
        )

def testFetchFeedBlobZeroDim(self):
        data = np.empty(shape=(2, 0, 3), dtype=np.float32)
        self.assertEqual(workspace.FeedBlob("testblob_empty", data), True)
        fetched_back = workspace.FetchBlob("testblob_empty")
        self.assertEqual(fetched_back.shape, (2, 0, 3))
        self.assertEqual(fetched_back.dtype, np.float32)

def testFetchFeedLongStringTensor(self):
        # long strings trigger array of object creation
        strs = np.array(
            [
                b" ".join(10 * [b"long string"]),
                b" ".join(128 * [b"very long string"]),
                b"small \0\1\2 string",
                b"Hello, world! I have special \0 symbols \1!",
            ]
        )
        workspace.FeedBlob("my_str_tensor", strs)
        strs2 = workspace.FetchBlob("my_str_tensor")
        self.assertEqual(strs.shape, strs2.shape)
        for i in range(0, strs.shape[0]):
            self.assertEqual(strs[i], strs2[i])

def testFetchFeedShortStringTensor(self):
        # small strings trigger NPY_STRING array
        strs = np.array([b"elem1", b"elem 2", b"element 3"])
        workspace.FeedBlob("my_str_tensor_2", strs)
        strs2 = workspace.FetchBlob("my_str_tensor_2")
        self.assertEqual(strs.shape, strs2.shape)
        for i in range(0, strs.shape[0]):
            self.assertEqual(strs[i], strs2[i])

def testFetchFeedPlainString(self):
        # this is actual string, not a tensor of strings
        s = b"Hello, world! I have special \0 symbols \1!"
        workspace.FeedBlob("my_plain_string", s)
        s2 = workspace.FetchBlob("my_plain_string")
        self.assertEqual(s, s2)

def testFetchBlobs(self):
        s1 = b"test1"
        s2 = b"test2"
        workspace.FeedBlob("s1", s1)
        workspace.FeedBlob("s2", s2)
        fetch1, fetch2 = workspace.FetchBlobs(["s1", "s2"])
        self.assertEquals(s1, fetch1)
        self.assertEquals(s2, fetch2)

def testFetchFeedViaBlobDict(self):
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        fetched = workspace.blobs["testblob"]
        # check if fetched is correct.
        self.assertEqual(fetched.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched, 1.0)
        fetched[:] = 2.0
        workspace.blobs["testblob"] = fetched
        fetched_again = workspace.blobs["testblob"]
        self.assertEqual(fetched_again.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched_again, 2.0)

self.assertTrue("testblob" in workspace.blobs)
        self.assertFalse("non_existant" in workspace.blobs)
        self.assertEqual(len(workspace.blobs), 1)
        for key in workspace.blobs:
            self.assertEqual(key, "testblob")

def testTorchInterop(self):
        workspace.RunOperatorOnce(
            core.CreateOperator(
                "ConstantFill", [], "foo", shape=(4,), value=2, dtype=10
            )
        )
        t = workspace.FetchTorch("foo")
        t.resize_(5)
        t[4] = t[2] = 777
        np.testing.assert_array_equal(t.numpy(), np.array([2, 2, 777, 2, 777]))
        np.testing.assert_array_equal(
            workspace.FetchBlob("foo"), np.array([2, 2, 777, 2, 777])
        )

z = torch.ones((4,), dtype=torch.int64)
        workspace.FeedBlob("bar", z)
        workspace.RunOperatorOnce(
            core.CreateOperator("Reshape", ["bar"], ["bar", "_"], shape=(2, 2))
        )
        z[0, 1] = 123
        np.testing.assert_array_equal(
            workspace.FetchBlob("bar"), np.array([[1, 123], [1, 1]])
        )
        np.testing.assert_array_equal(z, np.array([[1, 123], [1, 1]]))

class TestMultiWorkspaces(unittest.TestCase):
    def setUp(self):
        workspace.SwitchWorkspace("default")
        workspace.ResetWorkspace()

def testCreateWorkspace(self):
        self.net = core.Net("test-net")
        self.net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0)
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        self.assertEqual(workspace.HasBlob("testblob"), True)
        self.assertEqual(workspace.SwitchWorkspace("test", True), None)
        self.assertEqual(workspace.HasBlob("testblob"), False)
        self.assertEqual(workspace.SwitchWorkspace("default"), None)
        self.assertEqual(workspace.HasBlob("testblob"), True)

try:
            # The following should raise an error.
            workspace.SwitchWorkspace("non-existing")
            # so this should never happen.
            self.assertEqual(True, False)
        except RuntimeError:
            pass

workspaces = workspace.Workspaces()
        self.assertTrue("default" in workspaces)
        self.assertTrue("test" in workspaces)

@unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
class TestWorkspaceGPU(test_util.TestCase):
    def setUp(self):
        workspace.ResetWorkspace()
        self.net = core.Net("test-net")
        self.net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0)
        self.net.RunAllOnGPU()

def testFetchBlobGPU(self):
        self.assertEqual(
            workspace.RunNetOnce(self.net.Proto().SerializeToString()), True
        )
        fetched = workspace.FetchBlob("testblob")
        # check if fetched is correct.
        self.assertEqual(fetched.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched, 1.0)
        fetched[:] = 2.0
        self.assertEqual(workspace.FeedBlob("testblob", fetched), True)
        fetched_again = workspace.FetchBlob("testblob")
        self.assertEqual(fetched_again.shape, (1, 2, 3, 4))
        np.testing.assert_array_equal(fetched_again, 2.0)

def testGetGpuPeerAccessPattern(self):
        pattern = workspace.GetGpuPeerAccessPattern()
        self.assertEqual(type(pattern), np.ndarray)
        self.assertEqual(pattern.ndim, 2)
        self.assertEqual(pattern.shape[0], pattern.shape[1])
        self.assertEqual(pattern.shape[0], workspace.NumGpuDevices())

@unittest.skipIf(
        not workspace.has_cuda_support, "Tensor interop doesn't yet work on ROCm"
    )
    def testTorchInterop(self):
        # CUDA has convenient mem stats, let's use them to make sure we didn't
        # leak memory
        initial_mem = torch.cuda.memory_allocated()
        workspace.RunOperatorOnce(
            core.CreateOperator(
                "ConstantFill",
                [],
                "foo",
                shape=(4,),
                value=2,
                dtype=10,
                device_option=core.DeviceOption(workspace.GpuDeviceType),
            )
        )
        t = workspace.FetchTorch("foo")
        t.resize_(5)
        self.assertTrue(t.is_cuda)
        t[4] = t[2] = 777
        np.testing.assert_array_equal(t.cpu().numpy(), np.array([2, 2, 777, 2, 777]))
        np.testing.assert_array_equal(
            workspace.FetchBlob("foo"), np.array([2, 2, 777, 2, 777])
        )

z = torch.ones((4,), dtype=torch.int64, device="cuda")
        workspace.FeedBlob("bar", z)
        workspace.RunOperatorOnce(
            core.CreateOperator(
                "Reshape",
                ["bar"],
                ["bar", "_"],
                shape=(2, 2),
                device_option=core.DeviceOption(workspace.GpuDeviceType),
            )
        )
        z[0, 1] = 123
        np.testing.assert_array_equal(
            workspace.FetchBlob("bar"), np.array([[1, 123], [1, 1]])
        )
        np.testing.assert_array_equal(z.cpu(), np.array([[1, 123], [1, 1]]))

self.assertGreater(torch.cuda.memory_allocated(), initial_mem)
        # clean up everything
        del t
        del z
        workspace.ResetWorkspace()
        self.assertEqual(torch.cuda.memory_allocated(), initial_mem)

@unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.")
class TestWorkspaceIDEEP(test_util.TestCase):
    def testFeedFetchBlobIDEEP(self):
        arr = np.random.randn(2, 3).astype(np.float32)
        workspace.FeedBlob("testblob_ideep", arr, core.DeviceOption(caffe2_pb2.IDEEP))
        fetched = workspace.FetchBlob("testblob_ideep")
        np.testing.assert_array_equal(arr, fetched)

class TestImmedibate(test_util.TestCase):
    def testImmediateEnterExit(self):
        workspace.StartImmediate(i_know=True)
        self.assertTrue(workspace.IsImmediate())
        workspace.StopImmediate()
        self.assertFalse(workspace.IsImmediate())

def testImmediateRunsCorrectly(self):
        workspace.StartImmediate(i_know=True)
        net = core.Net("test-net")
        net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0)
        self.assertEqual(workspace.ImmediateBlobs(), ["testblob"])
        content = workspace.FetchImmediate("testblob")
        # Also, the immediate mode should not invade the original namespace,
        # so we check if this is so.
        with self.assertRaises(RuntimeError):
            workspace.FetchBlob("testblob")
        np.testing.assert_array_equal(content, 1.0)
        content[:] = 2.0
        self.assertTrue(workspace.FeedImmediate("testblob", content))
        np.testing.assert_array_equal(workspace.FetchImmediate("testblob"), 2.0)
        workspace.StopImmediate()
        with self.assertRaises(RuntimeError):
            content = workspace.FetchImmediate("testblob")

def testImmediateRootFolder(self):
        workspace.StartImmediate(i_know=True)
        # for testing we will look into the _immediate_root_folder variable
        # but in normal usage you should not access that.
        self.assertTrue(len(workspace._immediate_root_folder) > 0)
        root_folder = workspace._immediate_root_folder
        self.assertTrue(os.path.isdir(root_folder))
        workspace.StopImmediate()
        self.assertTrue(len(workspace._immediate_root_folder) == 0)
        # After termination, immediate mode should have the root folder
        # deleted.
        self.assertFalse(os.path.exists(root_folder))

class TestCppEnforceAsException(test_util.TestCase):
    def testEnforce(self):
        op = core.CreateOperator("Relu", ["X"], ["Y"])
        with self.assertRaises(RuntimeError):
            workspace.RunOperatorOnce(op)

class TestCWorkspace(htu.HypothesisTestCase):
    def test_net_execution(self):
        ws = workspace.C.Workspace()
        self.assertEqual(ws.nets, {})
        self.assertEqual(ws.blobs, {})
        net = core.Net("test-net")
        net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0)
        ws.create_net(net)
        # If we do not specify overwrite, this should raise an error.
        with self.assertRaises(RuntimeError):
            ws.create_net(net)
        # But, if we specify overwrite, this should pass.
        ws.create_net(net, True)
        # Overwrite can also be a kwarg.
        ws.create_net(net, overwrite=True)
        self.assertIn("testblob", ws.blobs)
        self.assertEqual(len(ws.nets), 1)
        net_name = net.Proto().name
        self.assertIn("test-net", net_name)
        net = ws.nets[net_name].run()
        blob = ws.blobs["testblob"]
        np.testing.assert_array_equal(
            np.ones((1, 2, 3, 4), dtype=np.float32), blob.fetch()
        )

@given(name=st.text(), value=st.floats(min_value=-1, max_value=1.0))
    def test_operator_run(self, name, value):
        ws = workspace.C.Workspace()
        op = core.CreateOperator("ConstantFill", [], [name], shape=[1], value=value)
        ws.run(op)
        self.assertIn(name, ws.blobs)
        np.testing.assert_allclose(
            [value], ws.blobs[name].fetch(), atol=1e-4, rtol=1e-4
        )

@given(
        blob_name=st.text(),
        net_name=st.text(),
        value=st.floats(min_value=-1, max_value=1.0),
    )
    def test_net_run(self, blob_name, net_name, value):
        ws = workspace.C.Workspace()
        net = core.Net(net_name)
        net.ConstantFill([], [blob_name], shape=[1], value=value)
        ws.run(net)
        self.assertIn(blob_name, ws.blobs)
        self.assertNotIn(net_name, ws.nets)
        np.testing.assert_allclose(
            [value], ws.blobs[blob_name].fetch(), atol=1e-4, rtol=1e-4
        )

@given(
        blob_name=st.text(),
        net_name=st.text(),
        plan_name=st.text(),
        value=st.floats(min_value=-1, max_value=1.0),
    )
    def test_plan_run(self, blob_name, plan_name, net_name, value):
        ws = workspace.C.Workspace()
        plan = core.Plan(plan_name)
        net = core.Net(net_name)
        net.ConstantFill([], [blob_name], shape=[1], value=value)

plan.AddStep(core.ExecutionStep("step", nets=[net], num_iter=1))

ws.run(plan)
        self.assertIn(blob_name, ws.blobs)
        self.assertIn(net.Name(), ws.nets)
        np.testing.assert_allclose(
            [value], ws.blobs[blob_name].fetch(), atol=1e-4, rtol=1e-4
        )

@given(
        blob_name=st.text(),
        net_name=st.text(),
        value=st.floats(min_value=-1, max_value=1.0),
    )
    def test_net_create(self, blob_name, net_name, value):
        ws = workspace.C.Workspace()
        net = core.Net(net_name)
        net.ConstantFill([], [blob_name], shape=[1], value=value)
        ws.create_net(net).run()
        self.assertIn(blob_name, ws.blobs)
        self.assertIn(net.Name(), ws.nets)
        np.testing.assert_allclose(
            [value], ws.blobs[blob_name].fetch(), atol=1e-4, rtol=1e-4
        )

@given(
        name=st.text(),
        value=htu.tensor(),
        device_option=st.sampled_from(htu.device_options),
    )
    def test_array_serde(self, name, value, device_option):
        ws = workspace.C.Workspace()
        ws.create_blob(name).feed(value, device_option=device_option)
        self.assertIn(name, ws.blobs)
        blob = ws.blobs[name]
        np.testing.assert_equal(value, ws.blobs[name].fetch())
        serde_blob = ws.create_blob("{}_serde".format(name))
        serde_blob.deserialize(blob.serialize(name))
        np.testing.assert_equal(value, serde_blob.fetch())

@given(name=st.text(), value=st.text())
    def test_string_serde(self, name, value):
        value = value.encode("ascii", "ignore")
        ws = workspace.C.Workspace()
        ws.create_blob(name).feed(value)
        self.assertIn(name, ws.blobs)
        blob = ws.blobs[name]
        self.assertEqual(value, ws.blobs[name].fetch())
        serde_blob = ws.create_blob("{}_serde".format(name))
        serde_blob.deserialize(blob.serialize(name))
        self.assertEqual(value, serde_blob.fetch())

def test_exception(self):
        ws = workspace.C.Workspace()

with self.assertRaises(TypeError):
            ws.create_net("...")

class TestPredictor(unittest.TestCase):
    def _create_model(self):
        m = model_helper.ModelHelper()
        y = brew.fc(
            m,
            "data",
            "y",
            dim_in=4,
            dim_out=2,
            weight_init=("ConstantFill", dict(value=1.0)),
            bias_init=("ConstantFill", dict(value=0.0)),
            axis=0,
        )
        m.net.AddExternalOutput(y)
        return m

# Use this test with a bigger model to see how using Predictor allows to
    # avoid issues with low protobuf size limit in Python
    #
    # def test_predictor_predefined(self):
    #     workspace.ResetWorkspace()
    #     path = 'caffe2/caffe2/test/assets/'
    #     with open(path + 'squeeze_predict_net.pb') as f:
    #         self.predict_net = f.read()
    #     with open(path + 'squeeze_init_net.pb') as f:
    #         self.init_net = f.read()
    #     self.predictor = workspace.Predictor(self.init_net, self.predict_net)

#     inputs = [np.zeros((1, 3, 256, 256), dtype='f')]
    #     outputs = self.predictor.run(inputs)
    #     self.assertEqual(len(outputs), 1)
    #     self.assertEqual(outputs[0].shape, (1, 1000, 1, 1))
    #     self.assertAlmostEqual(outputs[0][0][0][0][0], 5.19026289e-05)

def test_predictor_memory_model(self):
        workspace.ResetWorkspace()
        m = self._create_model()
        workspace.FeedBlob("data", np.zeros([4], dtype="float32"))
        self.predictor = workspace.Predictor(
            workspace.StringifyProto(m.param_init_net.Proto()),
            workspace.StringifyProto(m.net.Proto()),
        )

inputs = np.array([1, 3, 256, 256], dtype="float32")
        outputs = self.predictor.run([inputs])
        np.testing.assert_array_almost_equal(
            np.array([[516, 516]], dtype="float32"), outputs
        )

class TestTransform(htu.HypothesisTestCase):
    @given(
        input_dim=st.integers(min_value=1, max_value=10),
        output_dim=st.integers(min_value=1, max_value=10),
        batch_size=st.integers(min_value=1, max_value=10),
    )
    def test_simple_transform(self, input_dim, output_dim, batch_size):
        m = model_helper.ModelHelper()
        fc1 = brew.fc(m, "data", "fc1", dim_in=input_dim, dim_out=output_dim)
        fc2 = brew.fc(m, fc1, "fc2", dim_in=output_dim, dim_out=output_dim)
        conv = brew.conv(
            m,
            fc2,
            "conv",
            dim_in=output_dim,
            dim_out=output_dim,
            use_cudnn=True,
            engine="CUDNN",
            kernel=3,
        )

conv.Relu([], conv).Softmax([], "pred").LabelCrossEntropy(
            ["label"], ["xent"]
        ).AveragedLoss([], "loss")

transformed_net_proto = workspace.ApplyTransform("ConvToNNPack", m.net.Proto())

self.assertEqual(transformed_net_proto.op[2].engine, "NNPACK")

@given(
        input_dim=st.integers(min_value=1, max_value=10),
        output_dim=st.integers(min_value=1, max_value=10),
        batch_size=st.integers(min_value=1, max_value=10),
    )
    @settings(deadline=10000)
    def test_registry_invalid(self, input_dim, output_dim, batch_size):
        m = model_helper.ModelHelper()
        brew.fc(m, "data", "fc1", dim_in=input_dim, dim_out=output_dim)
        with self.assertRaises(RuntimeError):
            workspace.ApplyTransform("definitely_not_a_real_transform", m.net.Proto())

@given(value=st.floats(min_value=-1, max_value=1))
    @settings(deadline=10000)
    def test_apply_transform_if_faster(self, value):

init_net = core.Net("init_net")
        init_net.ConstantFill([], ["data"], shape=[5, 5, 5, 5], value=value)
        init_net.ConstantFill([], ["conv_w"], shape=[5, 5, 3, 3], value=value)
        init_net.ConstantFill([], ["conv_b"], shape=[5], value=value)

self.assertEqual(
            workspace.RunNetOnce(init_net.Proto().SerializeToString()), True
        )

m = model_helper.ModelHelper()
        conv = brew.conv(
            m,
            "data",
            "conv",
            dim_in=5,
            dim_out=5,
            kernel=3,
            use_cudnn=True,
            engine="CUDNN",
        )

conv.Relu([], conv).Softmax([], "pred").AveragedLoss([], "loss")

self.assertEqual(workspace.RunNetOnce(m.net.Proto().SerializeToString()), True)

proto = workspace.ApplyTransformIfFaster(
            "ConvToNNPack", m.net.Proto(), init_net.Proto()
        )
        self.assertEqual(workspace.RunNetOnce(proto.SerializeToString()), True)
        proto = workspace.ApplyTransformIfFaster(
            "ConvToNNPack",
            m.net.Proto(),
            init_net.Proto(),
            warmup_runs=10,
            main_runs=100,
            improvement_threshold=2.0,
        )
        self.assertEqual(workspace.RunNetOnce(proto.SerializeToString()), True)

class MyModule(torch.jit.ScriptModule):
    def __init__(self):
        super(MyModule, self).__init__()
        self.mult = torch.nn.Parameter(torch.tensor([[1, 2, 3, 4, 5.0]]))

@torch.jit.script_method
    def forward(self, x):
        return self.mult.mm(x)

@torch.jit.script_method
    def multi_input(self, x: torch.Tensor, y: torch.Tensor, z: int = 2) -> torch.Tensor:
        return x + y + z

@torch.jit.script_method
    def multi_input_tensor_list(self, tensor_list):  # pyre-ignore: PT type annotations
        # type: (List[Tensor]) -> Tensor
        return tensor_list[0] + tensor_list[1] + tensor_list[2]

@torch.jit.script_method
    def multi_output(self, x):
        return (x, x + 1)

@unittest.skipIf(
    "ScriptModule" not in core._REGISTERED_OPERATORS,
    "Script module integration in Caffe2 is not enabled",
)
class TestScriptModule(test_util.TestCase):
    def _createFeedModule(self):
        workspace.FeedBlob("m", MyModule())

def testCreation(self):
        m = MyModule()
        workspace.FeedBlob("module", m)
        m2 = workspace.FetchBlob("module")
        self.assertTrue(m2 is not None)

def testForward(self):
        self._createFeedModule()
        val = np.random.rand(5, 5).astype(np.float32)
        param = np.array([[1, 2, 3, 4, 5]]).astype(np.float32)
        workspace.FeedBlob("w", val)
        workspace.RunOperatorOnce(
            core.CreateOperator("ScriptModule", ["m", "w"], ["y"])
        )
        np.testing.assert_almost_equal(
            workspace.FetchBlob("y"), np.matmul(param, val), decimal=5
        )

def testMultiInputOutput(self):
        self._createFeedModule()
        val = np.random.rand(5, 5).astype(np.float32)
        workspace.FeedBlob("w", val)
        val2 = np.random.rand(5, 5).astype(np.float32)
        workspace.FeedBlob("w2", val2)
        workspace.RunOperatorOnce(
            core.CreateOperator(
                "ScriptModule", ["m", "w", "w2"], ["y"], method="multi_input"
            )
        )
        workspace.RunOperatorOnce(
            core.CreateOperator(
                "ScriptModule", ["m", "w"], ["y1", "y2"], method="multi_output"
            )
        )
        np.testing.assert_almost_equal(
            workspace.FetchBlob("y"), val + val2 + 2, decimal=5
        )
        np.testing.assert_almost_equal(workspace.FetchBlob("y1"), val, decimal=5)
        np.testing.assert_almost_equal(workspace.FetchBlob("y2"), val + 1, decimal=5)

def testMultiTensorListInput(self):
        self._createFeedModule()
        val = np.random.rand(5, 5).astype(np.float32)
        workspace.FeedBlob("w", val)
        val2 = np.random.rand(5, 5).astype(np.float32)
        workspace.FeedBlob("w2", val2)
        val3 = np.random.rand(5, 5).astype(np.float32)
        workspace.FeedBlob("w3", val3)

workspace.RunOperatorOnce(
            core.CreateOperator(
                "ScriptModule",
                ["m", "w", "w2", "w3"],
                ["y"],
                method="multi_input_tensor_list",
                pass_inputs_as_tensor_list=True,
            )
        )
        np.testing.assert_almost_equal(
            workspace.FetchBlob("y"), val + val2 + val3, decimal=5
        )

def testSerialization(self):
        tmpdir = tempfile.mkdtemp()
        try:
            self._createFeedModule()
            workspace.RunOperatorOnce(
                core.CreateOperator(
                    "Save",
                    ["m"],
                    [],
                    absolute_path=1,
                    db=os.path.join(tmpdir, "db"),
                    db_type="minidb",
                )
            )
            workspace.ResetWorkspace()

self.assertFalse(workspace.HasBlob("m"))
            workspace.RunOperatorOnce(
                core.CreateOperator(
                    "Load",
                    [],
                    [],
                    absolute_path=1,
                    db=os.path.join(tmpdir, "db"),
                    db_type="minidb",
                    load_all=1,
                )
            )
            self.assertTrue(workspace.HasBlob("m"))
            # TODO: make caffe2 side load return python-sided module
            # right now it returns the base class (torch._C.ScriptModule)
            # self.assertTrue(isinstance(workspace.FetchBlob('m'), torch.jit.ScriptModule))

# do something with the module
            val = np.random.rand(5, 5).astype(np.float32)
            param = np.array([[1, 2, 3, 4, 5]]).astype(np.float32)
            workspace.FeedBlob("w", val)
            workspace.RunOperatorOnce(
                core.CreateOperator("ScriptModule", ["m", "w"], ["y"])
            )
            np.testing.assert_almost_equal(
                workspace.FetchBlob("y"), np.matmul(param, val), decimal=5
            )
        finally:
            # clean up temp folder.
            try:
                shutil.rmtree(tmpdir)
            except OSError as e:
                if e.errno != errno.ENOENT:
                    raise

class TestScriptModuleFromString(TestScriptModule):
    def _createFeedModule(self):
        workspace.RunOperatorOnce(
            core.CreateOperator(
                "ScriptModuleLoad",
                [],
                ["m"],
                serialized_binary=self._get_modules_bytes(MyModule()),
            )
        )

def _get_modules_bytes(self, the_module):
        import io

buffer = io.BytesIO()
        torch.jit.save(the_module, buffer)
        return buffer.getvalue()

if __name__ == "__main__":
    unittest.main()

cmd: