<commit_before>"""The main drudge and tensor class definition."""

import functools
import inspect
import operator
import types
import typing
from collections.abc import Iterable, Sequence

from IPython.display import Math
from pyspark import RDD, SparkContext
from sympy import IndexedBase, Symbol, Indexed, Integer, Wild, latex

from .canonpy import Perm, Group
from .term import (
    Range, sum_term, Term, parse_term, Vec, subst_factor_in_term,
    subst_vec_in_term, parse_terms, einst_term
)
from .utils import ensure_symb, BCastVar, nest_bind, prod_


class Tensor:
    """The main tensor class.

    A tensor is an aggregate of terms distributed and managed by Spark.  Here
    most operations needed for tensors are defined.

    Normally, tensor instances are created from drudge methods or tensor
    operations.  Direct invocation of its constructor is seldom in user scripts.

    """

    __slots__ = [
        '_drudge',
        '_terms',
        '_local_terms',
        '_free_vars'
    ]

    #
    # Term creation
    #

    def __init__(self, drudge: 'Drudge', terms: RDD,
                 free_vars: typing.Set[Symbol] = None):
        """Initialize the tensor.

        This function is not designed to be called by users directly.  Tensor
        creation should be carried out by factory function inside drudges and
        the operations defined here.
        """

        self._drudge = drudge
        self._terms = terms

        self._local_terms = None  # type: typing.List[Term]
        self._free_vars = free_vars

    #
    # Basic information
    #

    @property
    def terms(self):
        """Get the terms as an RDD object.

        Although for users, normally there is no need for direct manipulation of
        the terms, it is still exposed here for flexibility.
        """
        return self._terms

    @property
    def local_terms(self) -> typing.List[Term]:
        """Gather the terms locally into a list.

        The list returned by this is for read-only and should **never** be
        mutated.<commit_msg>Add expansion and load-balancing info to tensor<commit_after>"""The main drudge and tensor class definition."""

import functools
import inspect
import operator
import types
import typing
from collections.abc import Iterable, Sequence

from IPython.display import Math
from pyspark import RDD, SparkContext
from sympy import IndexedBase, Symbol, Indexed, Integer, Wild, latex

from .canonpy import Perm, Group
from .term import (
    Range, sum_term, Term, parse_term, Vec, subst_factor_in_term,
    subst_vec_in_term, parse_terms, einst_term
)
from .utils import ensure_symb, BCastVar, nest_bind, prod_


class Tensor:
    """The main tensor class.

    A tensor is an aggregate of terms distributed and managed by Spark.  Here
    most operations needed for tensors are defined.

    Normally, tensor instances are created from drudge methods or tensor
    operations.  Direct invocation of its constructor is seldom in user scripts.

    """

    __slots__ = [
        '_drudge',
        '_terms',
        '_local_terms',
        '_free_vars',
        '_expanded',
        '_repartitioned'
    ]

    #
    # Term creation
    #

    def __init__(self, drudge: 'Drudge', terms: RDD,
                 free_vars: typing.Set[Symbol] = None,
                 expanded=False, repartitioned=False):
        """Initialize the tensor.

        This function is not designed to be called by users directly.  Tensor
        creation should be carried out by factory function inside drudges and
        the operations defined here.

        The default values for the keyword arguments are always the safest
        choice, for better performance, manipulations are encouraged to have
        proper consideration of all the keyword arguments.
        """

        self._drudge = drudge
        self._terms = terms

        self._local_terms = None  # type: typing.List[Term]

        self._free_vars = free_vars
        self._expanded = expanded
        self._repartitioned = repartitioned


    #
    # Basic information
    #

    @property
    def terms(self):
        """Get the terms as an RDD object.

        Although for users, normally there is no need for direct manipulation of
        the terms, it is still exposed here for flexibility.
        """
        return self._terms

    @property
    def local_terms(self) -> typing.List[Term]:
        """Gather the terms locally into a list.

        The list returned by this is for read-only and should **never** be
        mutated.
{'fraction_non_alphanumeric': 0.06556655665566556, 'fraction_numerical': 0.0, 'mean_word_length': 2.6194267515923566, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '299698', 'n_tokens_mistral': 1256, 'n_tokens_neox': 1184, 'n_words': 492}