Generates two sets of contraction axes for the two tensor arguments.
(a, axes)
| 4028 | return reshaped_a, free_dims, free_dims_static |
| 4029 | |
| 4030 | def _tensordot_axes(a, axes): |
| 4031 | """Generates two sets of contraction axes for the two tensor arguments.""" |
| 4032 | a_shape = a.get_shape() |
| 4033 | if isinstance(axes, compat.integral_types): |
| 4034 | if axes < 0: |
| 4035 | raise ValueError("'axes' must be at least 0.") |
| 4036 | if a_shape.ndims is not None: |
| 4037 | if axes > a_shape.ndims: |
| 4038 | raise ValueError("'axes' must not be larger than the number of " |
| 4039 | "dimensions of tensor %s." % a) |
| 4040 | return (list(xrange(a_shape.ndims - axes, |
| 4041 | a_shape.ndims)), list(xrange(axes))) |
| 4042 | else: |
| 4043 | rank = array_ops.rank(a) |
| 4044 | return (range(rank - axes, rank, |
| 4045 | dtype=dtypes.int32), range(axes, dtype=dtypes.int32)) |
| 4046 | elif isinstance(axes, (list, tuple)): |
| 4047 | if len(axes) != 2: |
| 4048 | raise ValueError("'axes' must be an integer or have length 2.") |
| 4049 | a_axes = axes[0] |
| 4050 | b_axes = axes[1] |
| 4051 | if isinstance(a_axes, compat.integral_types) and \ |
| 4052 | isinstance(b_axes, compat.integral_types): |
| 4053 | a_axes = [a_axes] |
| 4054 | b_axes = [b_axes] |
| 4055 | if len(a_axes) != len(b_axes): |
| 4056 | raise ValueError( |
| 4057 | "Different number of contraction axes 'a' and 'b', %s != %s." % |
| 4058 | (len(a_axes), len(b_axes))) |
| 4059 | return a_axes, b_axes |
| 4060 | else: |
| 4061 | axes = ops.convert_to_tensor(axes, name="axes", dtype=dtypes.int32) |
| 4062 | return axes[0], axes[1] |
| 4063 | |
| 4064 | with ops.name_scope(name, "Tensordot", [a, b, axes]) as name: |
| 4065 | a = ops.convert_to_tensor(a, name="a") |