python source code of mpo

Project: TensorNetwork (GitHub Link)

TensorNetwork-master
- AUTHORS
- dev_tools
  - Dockerfile
- examples
  - sat
    - SATTutorial.ipynb
    - sat_tensornetwork_test.py
    - sat_tensornetwork.py
    - __init__.py
  - fft
    - fft_test.py
    - fft.py
    - __init__.py
  - __init__.py
  - wavefunctions
    - wavefunctions_test.py
    - evolution_example.py
    - wavefunctions.py
    - trotter.py
  - simple_mera
    - simple_mera_test.py
    - simple_mera.py
- LICENSE
- CONTRIBUTING.md
- tensornetwork
  - backends
    - abstract_backend.py
    - backend_test.py
    - shell
      - shell_backend_test.py
      - shell_backend.py
      - __init__.py
    - backend_factory.py
    - jax
      - jitted_functions_test.py
      - jitted_functions.py
      - __init__.py
      - jax_backend_test.py
      - jax_backend.py
    - pytorch
      - pytorch_tensornetwork_test.py
      - decompositions.py
      - pytorch_backend_test.py
      - __init__.py
      - decompositions_test.py
      - pytorch_backend.py
    - symmetric
      - symmetric_backend_test.py
      - symmetric_backend.py
      - decompositions.py
      - __init__.py
      - decompositions_test.py
    - numpy
      - numpy_backend.py
      - decompositions.py
      - numpy_backend_test.py
      - __init__.py
      - decompositions_test.py
    - __init__.py
    - tensorflow
      - tensorflow_backend.py
      - tensorflow_backend_test.py
      - decompositions.py
      - tensorflow_tensornetwork_test.py
      - __init__.py
      - decompositions_test.py
      - tensordot2.py
      - tensordot2_test.py
  - config.py
  - ops.py
  - version.py
  - ncon_interface.py
  - component_factory.py
  - network_components.py
  - backend_contextmanager.py
  - block_sparse
    - blocksparsetensor.py
    - index.py
    - utils_test.py
    - blocksparsetensor_test.py
    - linalg.py
    - charge_test.py
    - charge.py
    - linalg_test.py
    - __init__.py
    - utils.py
    - index_test.py
    - tensordot_test.py
  - contractors
    - opt_einsum_paths
      - path_calculation_test.py
      - path_contractors.py
      - __init__.py
      - utils.py
      - path_contractors_node_test.py
    - bucket_contractor.py
    - __init__.py
    - bucket_contractor_test.py
  - tn_keras
    - condenser.py
    - entangler.py
    - mpo_test.py
    - images
    - expander.py
    - conv2d_mpo.py
    - test_layer.py
    - colabs
      - TN_Keras.ipynb
    - mpo.py
    - __init__.py
    - README.md
    - test_conv_layer.py
    - dense.py
    - TN_Keras.ipynb
  - __init__.py
  - utils.py
  - tests
    - split_node_symmetric_test.py
    - network_operations_test.py
    - ncon_interface_test.py
    - network_components_free_symmetric_test.py
    - tensornetwork_symmetric_test.py
    - split_node_test.py
    - __init__.py
    - backend_contextmanager_test.py
    - network_operations_symmetric_test.py
    - network_test.py
    - network_components_free_test.py
    - tensornetwork_test.py
  - quantum
    - quantum_test.py
    - quantum.py
  - visualization
    - __init__.py
    - graphviz_test.py
    - graphviz.py
  - linalg
    - linalg.py
    - linalg_test.py
    - __init__.py
  - matrixproductstates
    - base_mps_test.py
    - infinite_mps.py
    - infinite_mps_test.py
    - dmrg.py
    - mpo_test.py
    - finite_mps.py
    - mpo.py
    - __init__.py
    - dmrg_test.py
    - finite_mps_test.py
    - base_mps.py
  - network_operations.py
- setup.py
- colabs
  - Tensor_Networks_in_Neural_Networks.ipynb
- .style.yapf
- .travis.yml
- README.md
- requirements.txt
- pylintrc
- .gitignore
- docs
  - ncon.rst
  - Makefile
  - make.bat
  - block_sparse_tutorial.rst
  - contractors.rst
  - copy_contract.rst
  - quantum_circuit.rst
  - tutorial.rst
  - node_spliting.rst
  - block_sparse.rst
  - edges.rst
  - stubs
    - tensornetwork.BlockSparseTensor.rst
    - tensornetwork.Index.rst
    - tensornetwork.contractors.bucket.rst
    - tensornetwork.U1Charge.rst
    - tensornetwork.BaseCharge.rst
    - tensornetwork.contractors.greedy.rst
    - tensornetwork.Z2Charge.rst
    - tensornetwork.contractors.auto.rst
    - tensornetwork.FiniteMPS.rst
    - tensornetwork.BaseNode.rst
    - tensornetwork.Node.rst
    - tensornetwork.CopyNode.rst
    - tensornetwork.contractors.custom.rst
    - tensornetwork.contractors.branch.rst
    - tensornetwork.ChargeArray.rst
    - tensornetwork.ZNCharge.rst
    - tensornetwork.InfiniteMPS.rst
    - tensornetwork.NodeCollection.rst
    - tensornetwork.contractors.optimal.rst
    - tensornetwork.Edge.rst
  - mps.rst
  - nodes.rst
  - index.rst
  - conf.py
  - _static
  - contributing.rst
  - network.rst
- requirements_travis.txt
- MANIFEST.in
- conftest.py

import itertools
from tensornetwork.tn_keras import mpo
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras import Input
import numpy as np
import pytest


@pytest.mark.parametrize('in_dim_base,dim1,dim2,num_nodes,bond_dim',
                         itertools.product([3, 4], [3, 4], [2, 5], [3, 4],
                                           [2, 3]))
def test_shape_sanity_check(in_dim_base, dim1, dim2, num_nodes, bond_dim):
  model = Sequential([
      Input(in_dim_base**num_nodes),
      mpo.DenseMPO(dim1**num_nodes, num_nodes=num_nodes, bond_dim=bond_dim),
      mpo.DenseMPO(dim2**num_nodes, num_nodes=num_nodes, bond_dim=bond_dim),
  ])
  # Hard code batch size.
  result = model.predict(np.ones((32, in_dim_base**num_nodes)))
  assert result.shape == (32, dim2**num_nodes)