parameters: (op)
return_statement: return [diag_shape.concatenate(diag_shape[-1])]
docstring: Shape function for array_ops.batch_matrix_diag.
docstring_summary: Shape function for array_ops.batch_matrix_diag.
docstring_tokens: Shape function for array_ops . batch_matrix_diag .
function: def _BatchMatrixDiagShape(op):
    """Shape function for array_ops.batch_matrix_diag."""
    diag_shape = op.inputs[0].get_shape().with_rank_at_least(1)
    return [diag_shape.concatenate(diag_shape[-1])]
function_tokens: def _BatchMatrixDiagShape ( op ) : diag_shape = op . inputs [ 0 ] . get_shape ( ) . with_rank_at_least ( 1 ) return [ diag_shape . concatenate ( diag_shape [ - 1 ] ) ]